ai-search 搜索代码实现 SSE 版本

本文档详细说明了基于 SSE（Server-Sent Events）的实时消息传输接口及其代码实现。文中首先介绍了 SSE 接口的请求、响应及消息格式，然后逐个解析了各个关键模块的代码逻辑，包括请求处理、消息分发、搜索结果的处理、以及调用 DeepSeek 模型生成回答的过程。读者可以通过本文档了解整个系统的流程、各模块之间的交互关系及具体实现细节。

1. SSE 接口说明

1.1 接口描述

  • 接口地址：/api/chat/sse
  • 请求格式：客户端以 JSON 格式发送请求，示例如下：
{
  "type": "message",
  "userId": "7130314803843072",
  "message": {
    "messageId": "71383373052067817",
    "chatId": "7138337258229764",
    "content": "什么事deepseek"
  },
  "files": [],
  "focusMode": "webSearch",
  "copilotEnabled": true,
  "optimizationMode": "speed",
  "history": []
}

1.2 响应 Header

当请求为 SSE 请求时，响应头会包含以下信息，确保数据以事件流形式传输且连接保持长连接状态：

Name	                Value
server	            t-io
date	              Sun, 23 Mar 2025 10:40:10 GMT
content-type	      text/event-stream;charset=utf-8
connection	        keep-alive

1.3 响应 Body —— SSE 消息流

在 SSE 消息流中，每个 JSON 消息都有一个 type 字段，用于区分消息的用途。主要类型包括：

  1. sources
    用于在对话初期返回与用户查询相关的数据源信息，包含标题、URL 和内容摘要等信息。

  2. message
    携带实际要展示给用户的回答或文本内容。由于 SSE 支持分段传输，多个 message 类型的数据最终将拼接成完整回复。

  3. reasoning
    用于传递系统内部的链式推理（chain-of-thought）信息，主要供调试或内部分析时使用，前端可以选择不显示这些内容。

  4. messageEnd
    表示当前一次 SSE 消息传输的结束，前端可据此停止加载动画或关闭当前连接。

示例 SSE 消息流

下面展示了一个 SSE 消息流的示例：

data:{"data":[{"metadata":{"title":"DeepSeek：这是一家什么公司？为何突然引发美国科技股暴跌？ - BBC News 中文","url":"https://www.bbc.com/zhongwen/articles/ce854jy1j5go/simp"},"pageContent":"2025年1月28日 ... DeepSeek是一家什么公司？开发成本仅600万美元的软件为何引发美国科技股暴跌？ ... 中国公司DeepSeek（深度求索）开发的人工智慧聊天机器人自今年1月在美国发布 ..."},{"metadata":{"title":"DeepSeek：這是一家什麼公司？為何突然引發美國科技股暴跌？ - BBC News 中文","url":"https://www.bbc.com/zhongwen/articles/ce854jy1j5go/trad"},"pageContent":"圖像來源，Getty Images · 中國公司DeepSeek（深度求索）開發的人工智慧聊天機器人自今年1月在美國發布後，迅速成為蘋果商店下載次數最多的免費應用程式。"},{"metadata":{"title":"DeepSeek让我审视，什么才是属于“人”的创作 | 人人都是产品经理","url":"https://www.woshipm.com/share/6176371.html"},"pageContent":"随着DeepSeek等生成式AI的崛起，创作者们面临着前所未有的挑战与反思。本文通过作者与DeepSeek的深度互动，探讨了AI在模仿人类创作时的优缺点，以及它对人类创作者的启发与影响。文章不仅分析了DeepSeek的写作风格和创作..."},{"metadata":{"title":"DeepSeek 的秘方是硅谷味儿的-品玩","url":"https://www.pingwest.com/a/301544"},"pageContent":"科技创新者的每日必读 · DeepSeek 的秘方是硅谷味儿的"},{"metadata":{"title":"DeepSeek在海内外彻底爆发，但它不需要被神话 | 界面新闻","url":"https://m.jiemian.com/article/12298832.html"},"pageContent":"图片来源：界面图库 · 界面新闻记者 | 伍洋宇"},{"metadata":{"title":"揭秘DeepSeek:一个更极致的中国技术理想主义故事 ｜36氪独家-36氪","url":"https://www.36kr.com/p/2872793466982535"},"pageContent":"2 weeks ago - 三筒洗衣机、扫地机器人捡袜子……家电品牌是真听劝啊 · 美团开始向前看"},{"metadata":{"title":"对话郑永年：DeepSeek带来的震撼与启示--国际--人民网","url":"http://world.people.com.cn/n1/2025/0207/c1002-40414069.html"},"pageContent":"（图源：网络） · 1、侠客岛：今年春节前后，中国AI公司深度求索（DeepSeek）的大模型DeepSeek-R1火爆全网,很多人被它展现出来的强大能力震撼,您有没有使用？"},{"metadata":{"title":"三问三答方式告诉你：DeepSeek这款中国AI，为什么让美西方破大防_腾讯新闻","url":"https://news.qq.com/rain/a/20250203A05ZPD00"},"pageContent":"2025-02-03 21:36发布于河南历史领域创作者 · 1.中国人工智能初创公司DeepSeek在除夕前日成功引起全球关注，直接导致美国科技股大跌，如英伟达、台积电、谷歌等。"},{"metadata":{"title":"一文读懂DeepSeek（深度求索）：有限计算资源下的顶尖大模型_九方智投","url":"https://www.9fzt.com/9fztgw_1_top/e6bc1829bb1fef3d9e4988af00ae1882.html"},"pageContent":"一文读懂DeepSeek（深度求索）：有限计算资源下的顶尖大模型 · DeepSeek（深度求索）引起硅谷“恐慌”，性能超越了美国的顶尖模型，且研发成本更低，使用的芯片性能也较弱。甚至有人猜测DeepSeek（深度求索）模型对美国..."},{"metadata":{"title":"DeepSeek系列：什么是DeepSeek，它又能做些什么呢？ - 知乎","url":"https://zhuanlan.zhihu.com/p/24146505741"},"pageContent":"DeepSeek的核心技术基于深度学习，这是一种模仿人类大脑工作方式的技术。 它通过分析大量的文本、图片、语音等数据，学会了如何理解语言、生成内容，甚至解决复杂的问题。"},{"metadata":{"title":"关于DeepSeek，看这一篇就够了：全面解析与实用指南","url":"https://blog.csdn.net/2401_85343303/article/details/145462807"},"pageContent":"DeepSeek（深度求索）是一款由国内团队开发的开源人工智能工具库，专注于提供高效易用的AI模型训练与推理能力。 它既包含预训练大语言模型（如 DeepSeek -R1系列），也提供配套工具链，助力开发者快速实现 AI 应用落地。"},{"metadata":{"title":"DeepSeek_百度百科","url":"https://baike.baidu.com/item/DeepSeek/65258669"},"pageContent":"DeepSeek，全称杭州深度求索人工智能基础技术研究有限公司。DeepSeek是一家创新型科技公司，成立于2023年7月17日，使用数据蒸馏技术，得到更为精炼、有用的数据。"},{"metadata":{"title":"DeepSeek是什么？Deepseek最全使用基本技巧：从新手到 ...","url":"https://zhuanlan.zhihu.com/p/22495801392"},"pageContent":"DeepSeek 是由中国公司\"深度求索\"开发的智能对话模型，专注于通用人工智能（AGI）领域的研发。 它基于深度学习技术和大规模数据训练，能够完成对话生成、信息分析、任务执行等多种功能。"},{"metadata":{"title":"一文读懂| 关于DeepSeek公司及其大模型 - 新浪财经","url":"https://finance.sina.com.cn/stock/usstock/c/2025-01-27/doc-inehmkuy7811061.shtml?finpagefr=p_108"},"pageContent":"2025年1月27日 ... DeepSeek是2023年成立的中国初创企业，创始人梁文锋。其AI模型性能不逊于全球最强聊天机器人，成本低很多。1月初发布app迅速登顶美国iPhone下载榜， ..."},{"metadata":{"title":"什么是deepseek？它有哪些功能？对我们的生活有什么现实 ...","url":"https://www.zhihu.com/question/11238174016"},"pageContent":"DeepSeek是一款由国内人工智能公司深度求索（杭州深度求索人工智能基础技术研究有限公司）研发的大型语言模型，它拥有强大的自然语言处理能力，并集成多种AI技术。"},{"metadata":{"title":"深度求索- 维基百科，自由的百科全书","url":"https://zh.wikipedia.org/zh-hans/%E6%B7%B1%E5%BA%A6%E6%B1%82%E7%B4%A2"},"pageContent":"深度求索（DeepSeek），全称杭州深度求索人工智能基础技术研究有限公司，是中华人民共和国的一家人工智能与大型语言模型公司。该公司的总部位于中国大陆浙江省杭州市（注册 ..."},{"metadata":{"title":"DeepSeek科普篇：Deepseek为何一夜爆火？ - 知乎","url":"https://zhuanlan.zhihu.com/p/20924125507"},"pageContent":"DeepSeek，到底是什么来头？它为什么能一夜之间爆红，甚至搅动全球AI江湖？是国产AI真正的逆袭，还是昙花一现的泡沫？这篇文章，将对DeepSeek做一个全面的科普。一、deepseek为何能一鸣惊人？① 硬碰硬!DeepSeek R1 vs ChatGPT，谁更能打？"},{"metadata":{"title":"什么是 DeepSeek？为什么它会颠覆人工智能领域？ - OFweek维科网","url":"https://m.ofweek.com/ai/2025-02/ART-201700-8470-30656335.html"},"pageContent":"2025年2月5日 ... DeepSeek的算力来自于名叫“幻方量化”一家大型中国投资基金所有，该公司的GPU 数量是2048块H800 的很多倍。根据Scale AI 首席执行官亚历山大·王(Alexandr ..."},{"metadata":{"title":"关于DeepSeek的前世今生，看这一篇就够了 - 腾讯云开发者 ...","url":"https://cloud.tencent.com/developer/news/2133073"},"pageContent":"DeepSeek是一款由杭州深度求索人工智能基础技术研究有限公司打造的人工智能模型，具备理解自然语言并生成高质量文本内容的能力。本文介绍了DeepSeek的前世今生、性能、应用和创始人梁文锋的背景，以及它对美国AI领域的挑战和影响。"},{"metadata":{"title":"Deepseek是什么？一文全面了解Deepseek的作用和使用技巧！","url":"https://boardmix.cn/article/what-is-deepseek/"},"pageContent":"1. Deepseek是什么. Deepseek是一款由深度求索公司开发的智能助手产品，基于其强大的AI技术，能够帮助用户完成多种任务。Deepseek的主要功能特点如下： ... Deepseek具备高效的 ..."},{"metadata":{"title":"什么是DeepSeek？如何入门DeepSeek？ - 腾讯云","url":"https://cloud.tencent.com/developer/article/2493949"},"pageContent":"DeepSeek（深度求索）是一款由国内团队开发的开源人工智能工具库，专注于提供高效易用的AI模型训练与推理能力。 它既包含预训练大语言模型（如DeepSeek-R1系列），也提供配套工具链，助力开发者快速实现AI应用落地。"},{"metadata":{"title":"DeepSeek颠覆了什么？——大模型“国产之光”破局的启示 - 新华网","url":"http://www.news.cn/tech/20250207/f5d5fb35d6fb48bbb220e3438b682195/c.html"},"pageContent":"2025年2月7日 ... DeepSeek颠覆了什么？——大模型“国产之光”破局的启示-\"而DeepSeek重构了大模型的“成本-性能”曲线，同时压缩了计算周期。\""},{"metadata":{"title":"DeepSeek 到底是什么？——深入解析这一新兴技术的核心 ...","url":"https://baijiahao.baidu.com/s?id=1822470672068240488"},"pageContent":"DeepSeek 是一种基于深度学习和数据挖掘技术的智能搜索与分析系统，其核心在于通过深度学习模型和自然语言处理技术，理解数据的语义并提供精准的决策支持。"},{"metadata":{"title":"DeepSeek的惊人崛起是人工智能的“Sputnik时刻”还是“珍珠港事件”","url":"https://www.bbc.com/zhongwen/articles/c2d3n28dnwlo/simp"},"pageContent":"2025年1月28日 ... 分析：DeepSeek的惊人崛起是人工智能的“Sputnik时刻”还是“珍珠港事件” ... 在美国总统特朗普重返白宫的第一个周末，中国科技界向矽谷及白宫高层投下了震撼消息 ..."},{"metadata":{"title":"DeepSeek是什么？能干什么？如何高效为我所用？ - 百家号","url":"https://baijiahao.baidu.com/s?id=1823641543280391003"},"pageContent":"DeepSeek，顾名思义，是一款深度求索的 AI 工具。 它集自然语言处理（ NLP ）、计算机视觉（ CV ）、语音识别等多个领域的先进技术于一体，为用户提供了高效、便捷的 AI 模型训练、部署和应用服务。"},{"metadata":{"title":"对话幻方梁文锋：AI界的拼多多——揭秘DeepSeek！ - 华尔街见闻","url":"https://wallstreetcn.com/articles/3719982"},"pageContent":"2024年7月17日 ... 在硅谷，DeepSeek被称作“来自东方的神秘力量”。与很多大厂烧钱补贴不同，DeepSeek是有利润的。 中国的7家大模型创业公司中，DeepSeek（ ..."},{"metadata":{"title":"爆火的DeepSeek到底是什么？（一次性讲解清楚） - CSDN博客","url":"https://blog.csdn.net/dhdjjfhdghh/article/details/145475205"},"pageContent":"DeepSeek模型是由DeepSeek团队开发的开源大规模语言模型系列，旨在探索通用人工智能（AGI）的奥秘，并以长期主义视角回答人工智能领域的关键问题。 其开发背景是当前人工智能领域对于更高效、更强大语言模型的需求，以及对开源模型在性能和成本效益上与闭 ..."},{"metadata":{"title":"银行人有了新“上班搭子”，DeepSeek会给银行带来什么？_银行_金融频道首页_财经网 - CAIJING.COM.CN","url":"http://finance.caijing.com.cn/20250205/5068227.shtml"},"pageContent":"今年春节，DeepSeek火爆全球，凭借高性能、开源的模型，DeepSeek在AI界与科技圈激起千层浪。在追求高效、精准决策的金融业，DeepSeek的融入也是恰逢其时，这位赛博同事上线后仅需一秒，就能生成一份专业且精准的理财方案..."},{"metadata":{"title":"关于DeepSeek，你可能还不知道的10件事 - 21经济网","url":"https://www.21jingji.com/article/20250206/herald/691391dd091a0516241df9656ab755da.html"},"pageContent":"2025年1月20日，中国初创企业深度求索（DeepSeek）推出开源大模型DeepSeek-R1。“极低成本对标顶尖性能”的技术突破，引得一众硅谷科技巨头连夜发声。短短一周内，DeepSeek登顶全球140个市场的应用商店下载榜，甚至因访问量..."},{"metadata":{"title":"事关DeepSeek，硅谷大佬“互掐”！OpenAI又爆巨额融资","url":"http://www.stcn.com/article/detail/1512982.html"},"pageContent":"DeepSeek在硅谷持续引发讨论。 · 中国大模型创业公司DeepSeek因此前发布的开源模型V3和R1，受到了全球的关注。由于成本低廉、性能良好、透明公开，DeepSeek让人们开始质疑AI投资可能存在的泡沫，也挑战了硅谷在人工智能领..."},{"metadata":{"title":"变天！——由一个惊艳例子引发对DeepSeek的10条思考 | 人人都是产品经理","url":"https://www.woshipm.com/aigc/6174864.html"},"pageContent":"DeepSeek R1模型以其卓越的性能和惊艳的输出，在AI领域引起了广泛关注。本文通过一个具体的案例——模拟李世民在玄武门之变后的独白，展示了DeepSeek R1在历史细节、文学表达和用户需求预判方面的强大能力。 · 如果我们..."},{"metadata":{"title":"什么是 DeepSeek？为什么它会颠覆人工智能领域？_腾讯新闻","url":"https://news.qq.com/rain/a/20250131A05NIP00"},"pageContent":"今年，中国农历蛇年的春节假期，对于 AI 科技界最火的两个东西，莫非和Open AI ChatGPT 类似的 Deepseek 以及春晚表演节目的人形机器人。其中Deepseek 推理模型R1的推出更是让美国总统川普说成“A wakeup call”，而且此模型推出..."},{"metadata":{"title":"浪潮云率先发布，事关DeepSeek！算力概念持续走高，浪潮信息、用友网络等多股涨停，寒武纪涨超7% | 每经网","url":"https://cd.nbd.com.cn/articles/2025-02-07/3744260.html"},"pageContent":"每经网首页 > 要闻 > 正文 · 每日经济新闻 2025-02-07 12:19:06"},{"metadata":{"title":"揭秘DeepSeek:一个更极致的中国技术理想主义故事_新浪财经_新浪网","url":"https://finance.sina.com.cn/stock/roll/2024-12-29/doc-inecavfz5743886.shtml?cre=tianyi&mod=pchp&loc=4&r=0&rfunc=1&tj=cxvertical_pc_hp&tr=12"},"pageContent":"炒股就看金麒麟分析师研报，权威，专业，及时，全面，助您挖掘潜力主题机会！ · 来源：暗涌Waves"},{"metadata":{"title":"华尔街见闻","url":"https://wallstreetcn.com/articles/3739862"},"pageContent":"We cannot provide a description for this page right now"},{"metadata":{"title":"突然大批涌现！有人几天收入超5万元，事关DeepSeek - 21经济网","url":"https://www.21jingji.com/article/20250206/herald/0a4e43c2075afc921d7af31c43ae677c.html"},"pageContent":"近日，国内领先的人工智能公司深度求索推出的AI智能助手DeepSeek爆火，成为全网热议的焦点。 · 目前，DeepSeek研发团队暂未对外公布营收数据。DeepSeek赚不赚钱还不清楚，但有人已经通过“DeepSeek培训课”赚钱。"},{"metadata":{"title":"欧洲媒体：DeepSeek证明，科技制裁是一条错误的路_新闻频道_中国青年网","url":"https://news.youth.cn/gj/202502/t20250205_15808834.htm"},"pageContent":"欧洲媒体：DeepSeek证明，科技制裁是一条错误的路 · 发稿时间：2025-02-05 15:26:00 作者：袁野 来源： 中国青年报客户端"},{"metadata":{"title":"All Sources","url":"//127.0.0.1/sources/71383373052067817"}}],"messageId":493612582016638976,"type":"sources"}

data:{"data":"","messageId":493612582016638976,"type":"message"}

data:{"data":"start thinking...","messageId":493612582016638976,"type":"message"}

data:{"data":"嗯","messageId":493612582016638976,"type":"reasoning"}

data:{"data":"，","messageId":493612582016638976,"type":"reasoning"}

data:{"data":"用户","messageId":493612582016638976,"type":"reasoning"}

data:{"data":"省略其他消息","messageId":493612582016638976,"type":"reasoning"}

data:{"data":"反思","messageId":493612582016638976,"type":"reasoning"}

data:{"data":"。\n","messageId":493612582016638976,"type":"reasoning"}

data:{"data":"\n\n","messageId":493612582016638976,"type":"message"}

data:{"data":"Deep","messageId":493612582016638976,"type":"message"}

data:{"data":"省略其他消息","messageId":493612582016638976,"type":"message"}

data:{"data":"]","messageId":493612582016638976,"type":"message"}

data:{"messageId":493612582016638976,"type":"messageEnd"}

2. 代码实现详解

下面依次介绍各个模块的代码实现和业务逻辑：

2.1 ChatSSEHandler

该类作为 SSE 请求的入口，负责设置响应头、解析请求并调用业务层进行消息分发，同时处理异常情况。

package com.litongjava.perplexica.handler;

import com.alibaba.fastjson2.JSONObject;
import com.litongjava.jfinal.aop.Aop;
import com.litongjava.perplexica.services.WsChatService;
import com.litongjava.perplexica.vo.ChatWsReqMessageVo;
import com.litongjava.perplexica.vo.ChatWsRespVo;
import com.litongjava.tio.boot.http.TioRequestContext;
import com.litongjava.tio.core.ChannelContext;
import com.litongjava.tio.core.Tio;
import com.litongjava.tio.http.common.HttpRequest;
import com.litongjava.tio.http.common.HttpResponse;
import com.litongjava.tio.http.common.sse.SsePacket;
import com.litongjava.tio.http.server.util.CORSUtils;
import com.litongjava.tio.http.server.util.SseEmitter;
import com.litongjava.tio.utils.json.FastJson2Utils;
import com.litongjava.tio.websocket.common.WebSocketResponse;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class ChatSSEHandler {
  public HttpResponse chat(HttpRequest request) {
    HttpResponse response = TioRequestContext.getResponse();
    CORSUtils.enableCORS(response);
    ChannelContext channelContext = request.getChannelContext();
    // 设置sse请求头
    response.addServerSentEventsHeader();
    // 手动发送消息到客户端,因为已经设置了sse的请求头,所以客户端的连接不会关闭
    Tio.bSend(channelContext, response);
    response.setSend(false);

    String text = request.getBodyString();
    JSONObject reqJsonObject = FastJson2Utils.parseObject(text);
    String type = reqJsonObject.getString("type");
    if ("message".equals(type)) {
      ChatWsReqMessageVo vo = FastJson2Utils.parse(text, ChatWsReqMessageVo.class);
      vo.setSse(true);
      log.info("message:{}", text);
      try {
        Aop.get(WsChatService.class).dispatch(channelContext, vo);
      } catch (Exception e) {
        log.error(e.getMessage(), e);
        if(vo.isSse()) {
          ChatWsRespVo<String> error = ChatWsRespVo.error(e.getClass().toGenericString(), e.getMessage());
          byte[] jsonBytes = FastJson2Utils.toJSONBytes(error);
          Tio.bSend(channelContext, new SsePacket(jsonBytes));
          SseEmitter.closeSeeConnection(channelContext);

        }else {
          ChatWsRespVo<String> error = ChatWsRespVo.error(e.getClass().toGenericString(), e.getMessage());
          WebSocketResponse packet = WebSocketResponse.fromJson(error);
          Tio.bSend(channelContext, packet);

        }
      }
    }
    return response;
  }

}

说明

  • 设置 SSE 请求头：调用 response.addServerSentEventsHeader() 为响应设置 SSE 必要的 HTTP 头。
  • 保持连接：通过 Tio.bSend 发送响应并将 response.setSend(false)，使连接保持开启状态。
  • 消息解析与分发：将请求体转换为 ChatWsReqMessageVo 对象，设置 SSE 标记为 true 后调用 WsChatService.dispatch 方法进行消息处理。
  • 异常处理：如果处理过程中出现异常，根据请求是否为 SSE 分别发送错误消息并关闭 SSE 连接。

2.2 WsChatService

该类为业务层的核心，主要负责处理用户消息、创建或查询会话、问题重写、消息存储以及根据不同的 focusMode 调用相应的处理逻辑（如搜索、翻译、DeepSeek 等）。

package com.litongjava.perplexica.services;

import java.time.Instant;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.locks.Lock;

import com.google.common.util.concurrent.Striped;
import com.jfinal.kit.Kv;
import com.litongjava.db.activerecord.Db;
import com.litongjava.gemini.GoogleGeminiModels;
import com.litongjava.google.search.GoogleCustomSearchResponse;
import com.litongjava.google.search.SearchResultItem;
import com.litongjava.jfinal.aop.Aop;
import com.litongjava.model.web.WebPageContent;
import com.litongjava.openai.chat.ChatMessage;
import com.litongjava.openai.chat.OpenAiChatMessage;
import com.litongjava.openai.chat.OpenAiChatRequestVo;
import com.litongjava.openai.client.OpenAiClient;
import com.litongjava.openai.constants.PerplexityConstants;
import com.litongjava.openai.constants.PerplexityModels;
import com.litongjava.perplexica.callback.PerplexiticySeeCallback;
import com.litongjava.perplexica.callback.SearchGeminiSseCallback;
import com.litongjava.perplexica.can.ChatWsStreamCallCan;
import com.litongjava.perplexica.consts.FocusMode;
import com.litongjava.perplexica.consts.SearchTableNames;
import com.litongjava.perplexica.model.MaxSearchChatMessage;
import com.litongjava.perplexica.model.MaxSearchChatSession;
import com.litongjava.perplexica.vo.ChatParamVo;
import com.litongjava.perplexica.vo.ChatReqMessage;
import com.litongjava.perplexica.vo.ChatWsReqMessageVo;
import com.litongjava.perplexica.vo.ChatWsRespVo;
import com.litongjava.perplexica.vo.CitationsVo;
import com.litongjava.perplexica.vo.WebPageSource;
import com.litongjava.template.PromptEngine;
import com.litongjava.tio.core.ChannelContext;
import com.litongjava.tio.core.Tio;
import com.litongjava.tio.http.common.sse.SsePacket;
import com.litongjava.tio.utils.environment.EnvUtils;
import com.litongjava.tio.utils.json.FastJson2Utils;
import com.litongjava.tio.utils.json.JsonUtils;
import com.litongjava.tio.utils.snowflake.SnowflakeIdUtils;
import com.litongjava.tio.utils.tag.TagUtils;
import com.litongjava.tio.utils.thread.TioThreadUtils;
import com.litongjava.tio.websocket.common.WebSocketResponse;

import lombok.extern.slf4j.Slf4j;
import okhttp3.Call;
import okhttp3.Callback;

@Slf4j
public class WsChatService {
  private static final Striped<Lock> sessionLocks = Striped.lock(1024);
  private GeminiPredictService geminiPredictService = Aop.get(GeminiPredictService.class);
  private AiSearchService aiSerchService = Aop.get(AiSearchService.class);
  private SummaryQuestionService summaryQuestionService = Aop.get(SummaryQuestionService.class);
  private ChatMessgeService chatMessgeService = Aop.get(ChatMessgeService.class);
  private WebpageSourceService webpageSourceService = Aop.get(WebpageSourceService.class);

  /**
   * 使用搜索模型处理消息
  */
  public void dispatch(ChannelContext channelContext, ChatWsReqMessageVo reqMessageVo) {
    ChatReqMessage message = reqMessageVo.getMessage();
    Long userId = reqMessageVo.getUserId();
    Long sessionId = message.getChatId();
    Long messageQuestionId = message.getMessageId();
    String content = message.getContent();

    ChatParamVo chatParamVo = new ChatParamVo();
    // create chat or save message
    String focusMode = reqMessageVo.getFocusMode();
    if (!Db.exists(SearchTableNames.max_search_chat_session, "id", sessionId)) {
      Lock lock = sessionLocks.get(sessionId);
      lock.lock();
      try {
        TioThreadUtils.execute(() -> {
          String summary = summaryQuestionService.summary(content);
          new MaxSearchChatSession().setId(sessionId).setUserId(userId).setTitle(summary).setFocusMode(focusMode).save();
        });
      } finally {
        lock.unlock();
      }
    }

    // query history
    List<ChatMessage> history = chatMessgeService.getHistoryById(sessionId);
    chatParamVo.setHistory(history);

    if (content.length() > 30 || history.size() > 0) {
      String rewrited = Aop.get(RewriteQuestionService.class).rewrite(content, history);
      log.info("rewrite to:{}", rewrited);
      chatParamVo.setRewrited(rewrited);
      if (channelContext != null) {
        Kv end = Kv.by("type", "rewrited").set("content", rewrited);
        byte[] jsonBytes = FastJson2Utils.toJSONBytes(end);
        if (reqMessageVo.isSse()) {
          Tio.bSend(channelContext, new SsePacket(jsonBytes));
        } else {
          Tio.bSend(channelContext, WebSocketResponse.fromBytes(jsonBytes));
        }
      }
    }

    // save user mesasge
    new MaxSearchChatMessage().setId(messageQuestionId).setChatId(sessionId)
        //
        .setRole("user").setContent(content).save();

    String from = channelContext.getString("FROM");
    chatParamVo.setFrom(from);

    Boolean copilotEnabled = reqMessageVo.getCopilotEnabled();
    Call call = null;
    long answerMessageId = SnowflakeIdUtils.id();
    chatParamVo.setAnswerMessageId(answerMessageId);

    log.info("focusMode:{},{}", userId, focusMode);
    if (FocusMode.webSearch.equals(focusMode)) {
      call = aiSerchService.search(channelContext, reqMessageVo, chatParamVo);

    } else if (FocusMode.translator.equals(focusMode)) {
      String inputPrompt = Aop.get(TranslatorPromptService.class).genInputPrompt(channelContext, content, copilotEnabled, messageQuestionId, messageQuestionId, from);
      chatParamVo.setSystemPrompt(inputPrompt);
      call = geminiPredictService.predict(channelContext, reqMessageVo, chatParamVo);

    } else if (FocusMode.deepSeek.equals(focusMode)) {
      Aop.get(DeepSeekPredictService.class).predict(channelContext, reqMessageVo, chatParamVo);

    } else if (FocusMode.mathAssistant.equals(focusMode)) {
      String inputPrompt = PromptEngine.renderToString("math_assistant_prompt.txt");
      chatParamVo.setSystemPrompt(inputPrompt);
      Aop.get(DeepSeekPredictService.class).predict(channelContext, reqMessageVo, chatParamVo);

    } else if (FocusMode.writingAssistant.equals(focusMode)) {
      String inputPrompt = PromptEngine.renderToString("writing_assistant_prompt.txt");
      chatParamVo.setSystemPrompt(inputPrompt);
      Aop.get(DeepSeekPredictService.class).predict(channelContext, reqMessageVo, chatParamVo);
    } else {
      // 5. 向前端通知一个空消息，标识搜索结束，开始推理
      //{"type":"message","data":"", "messageId": "32fcbbf251337c"}
      ChatWsRespVo<String> chatVo = ChatWsRespVo.message(answerMessageId, "");
      byte[] jsonBytes = FastJson2Utils.toJSONBytes(chatVo);

      if (channelContext != null) {
        if (reqMessageVo.isSse()) {
          Tio.bSend(channelContext, new SsePacket(jsonBytes));
        } else {
          Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
        }
      }

      chatVo = ChatWsRespVo.message(answerMessageId, "Sorry Developing");
      jsonBytes = FastJson2Utils.toJSONBytes(chatVo);
      if (channelContext != null) {
        if (reqMessageVo.isSse()) {
          Tio.bSend(channelContext, new SsePacket(jsonBytes));
        } else {
          Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
        }

        Kv end = Kv.by("type", "messageEnd").set("messageId", answerMessageId);
        jsonBytes = FastJson2Utils.toJSONBytes(end);

        if (reqMessageVo.isSse()) {
          Tio.bSend(channelContext, new SsePacket(jsonBytes));
        } else {
          Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
        }
      }
    }

    if (call != null) {
      ChatWsStreamCallCan.put(sessionId.toString(), call);
    }
  }

  public Call google(ChannelContext channelContext, Long sessionId, Long messageId, String content, boolean isSSE) {
    String cseId = (String) channelContext.getString("CSE_ID");

    long answerMessageId = SnowflakeIdUtils.id();
    //1.问题重写
    // 省略
    //2.搜索
    GoogleCustomSearchResponse search = Aop.get(GoogleCustomSearchService.class).search(cseId, content);
    List<SearchResultItem> items = search.getItems();
    List<WebPageContent> results = new ArrayList<>(items.size());
    for (SearchResultItem searchResultItem : items) {
      String title = searchResultItem.getTitle();
      String link = searchResultItem.getLink();
      String snippet = searchResultItem.getSnippet();
      WebPageContent searchSimpleResult = new WebPageContent(title, link, snippet);
      results.add(searchSimpleResult);
    }
    //3.选择
    Kv kv = Kv.by("quesiton", content).set("search_result", JsonUtils.toJson(results));
    String fileName = "WebSearchSelectPrompt.txt";
    String prompt = PromptEngine.renderToString(fileName, kv);
    log.info("WebSearchSelectPrompt:{}", prompt);

    String selectResultContent = Aop.get(GeminiService.class).generate(prompt);
    List<String> outputs = TagUtils.extractOutput(selectResultContent);
    String titleAndLinks = outputs.get(0);
    if ("not_found".equals(titleAndLinks)) {
      //{"type":"message","data":"", "messageId": "32fcbbf251337c"}

      if (channelContext != null) {
        ChatWsRespVo<String> vo = ChatWsRespVo.message(answerMessageId, "");
        Tio.bSend(channelContext, WebSocketResponse.fromJson(vo));
        vo = ChatWsRespVo.message(messageId, "Sorry,not found");
        log.info("not found:{}", content);
        Tio.bSend(channelContext, WebSocketResponse.fromJson(vo));
      }

      return null;
    }
    //4.send to client
    String[] split = titleAndLinks.split("\n");
    List<CitationsVo> citationList = new ArrayList<>();
    for (int i = 0; i < split.length; i++) {
      String[] split2 = split[i].split("~~");
      citationList.add(new CitationsVo(split2[0], split2[1]));
    }

    if (citationList.size() > 0) {

      List<WebPageSource> sources = webpageSourceService.getListWithCitationsVo(citationList);
      ChatWsRespVo<List<WebPageSource>> chatRespVo = new ChatWsRespVo<>();
      chatRespVo.setType("sources").setData(sources).setMessageId(answerMessageId);
      byte[] jsonBytes = FastJson2Utils.toJSONBytes(chatRespVo);
      if (channelContext != null) {
        if (isSSE) {
          Tio.bSend(channelContext, new SsePacket(jsonBytes));
        } else {
          Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
        }
      }
    }

    //{"type":"message","data":"", "messageId": "32fcbbf251337c"}
    ChatWsRespVo<String> vo = ChatWsRespVo.message(answerMessageId, "");
    byte[] jsonBytes = FastJson2Utils.toJSONBytes(vo);
    if (channelContext != null) {
      if (isSSE) {
        Tio.bSend(channelContext, new SsePacket(jsonBytes));
      } else {
        Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));

      }
    }

    StringBuffer pageContents = Aop.get(SpiderService.class).spiderAsync(channelContext, answerMessageId, citationList);
    //6.推理
    String isoTimeStr = DateTimeFormatter.ISO_INSTANT.format(Instant.now());

    kv = Kv.by("date", isoTimeStr).set("context", pageContents.toString());
    String webSearchResponsePrompt = PromptEngine.renderToString("WebSearchResponsePrompt.txt", kv);
    log.info("webSearchResponsePrompt:{}", webSearchResponsePrompt);

    List<OpenAiChatMessage> messages = new ArrayList<>();
    messages.add(new OpenAiChatMessage("assistant", webSearchResponsePrompt));
    messages.add(new OpenAiChatMessage(content));

    OpenAiChatRequestVo chatRequestVo = new OpenAiChatRequestVo().setModel(GoogleGeminiModels.GEMINI_2_0_FLASH_EXP)
        //
        .setMessages(messages).setMax_tokens(3000);
    chatRequestVo.setStream(true);
    long start = System.currentTimeMillis();

    Callback callback = new SearchGeminiSseCallback(channelContext, sessionId, messageId, answerMessageId, start);
    Call call = Aop.get(GeminiService.class).stream(chatRequestVo, callback);
    return call;
  }

}

说明

  • 会话管理：判断当前会话是否存在，若不存在则通过分布式锁创建新会话，并调用 summaryQuestionService 生成问题摘要。
  • 历史消息查询：查询会话历史记录，并在必要时进行问题重写，将重写后的问题通知前端。
  • 消息保存：将用户消息保存到数据库。
  • FocusMode 分支处理：根据 focusMode 不同，调用不同的处理服务，如 webSearch、translator、deepSeek、mathAssistant 或 writingAssistant。
  • google 方法：展示了基于 Google Custom Search 的搜索逻辑，包括搜索结果获取、选择、引用信息生成以及抓取页面内容，最终调用 GeminiService 的流式接口生成回答。

2.3 AiSearchService

该服务处理搜索请求，根据用户设置（如 copilotEnabled、optimizationMode）决定是否进行搜索，再调用 Tavily Search API 获取网页内容并生成提示词以供后续回答生成使用。

package com.litongjava.max.search.services;

import java.time.Instant;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;

import org.postgresql.util.PGobject;

import com.jfinal.kit.Kv;
import com.litongjava.db.activerecord.Db;
import com.litongjava.jfinal.aop.Aop;
import com.litongjava.kit.PgObjectUtils;
import com.litongjava.max.search.consts.OptimizationMode;
import com.litongjava.max.search.vo.ChatParamVo;
import com.litongjava.max.search.vo.ChatWsReqMessageVo;
import com.litongjava.max.search.vo.ChatWsRespVo;
import com.litongjava.max.search.vo.WebPageSource;
import com.litongjava.model.web.WebPageContent;
import com.litongjava.template.PromptEngine;
import com.litongjava.tio.core.ChannelContext;
import com.litongjava.tio.core.Tio;
import com.litongjava.tio.http.common.RequestHeaderKey;
import com.litongjava.tio.http.common.sse.SsePacket;
import com.litongjava.tio.utils.hutool.StrUtil;
import com.litongjava.tio.utils.json.FastJson2Utils;
import com.litongjava.tio.websocket.common.WebSocketResponse;

import lombok.extern.slf4j.Slf4j;
import okhttp3.Call;

@Slf4j
public class AiSearchService {
  public PredictService predictService = Aop.get(PredictService.class);
  private AiRankerService aiRankerService = Aop.get(AiRankerService.class);
  private MaxSearchSearchService maxSearchSearchService = Aop.get(MaxSearchSearchService.class);
  private VectorRankerService vectorRankerService = Aop.get(VectorRankerService.class);
  public boolean spped = true;

  /**
   * 处理搜索请求
   * 根据当前用户的设置（copilotEnabled、optimizationMode）决定是否进行搜索，
   * 然后调用 Tavily Search（通过 MaxSearchSearchService）获取网页内容，
   * 并进一步生成提示词供后续回答生成使用。
   *
   * @param channelContext 通道上下文，用于返回消息
   * @param reqMessageVo   用户请求消息对象，包含消息内容、用户设置及历史记录
   * @param chatParamVo    对话参数对象，保存问题重写结果、提示词、搜索结果等信息
   * @return 返回用于流式处理回答的 Call 对象
   */
  public Call search(ChannelContext channelContext, ChatWsReqMessageVo reqMessageVo, ChatParamVo chatParamVo) {
    String optimizationMode = reqMessageVo.getOptimizationMode();
    Boolean copilotEnabled = reqMessageVo.getCopilotEnabled();
    String content = reqMessageVo.getMessage().getContent();
    Long questionMessageId = reqMessageVo.getMessage().getMessageId();
    long answerMessageId = chatParamVo.getAnswerMessageId();

    String inputPrompt = null;
    if (copilotEnabled != null && copilotEnabled) {
      String quesiton = null;
      // 如果有问题重写，则优先使用重写后的问题，否则直接使用原始内容
      if (chatParamVo.getRewrited() != null) {
        quesiton = chatParamVo.getRewrited();
      } else {
        quesiton = content;
      }

      // 使用 MaxSearchSearchService 对 Tavily Search API 进行调用
      List<WebPageContent> webPageContents = maxSearchSearchService.search(quesiton);

      // 根据优化模式对搜索结果进行处理
      JinaReaderService jinaReaderService = Aop.get(JinaReaderService.class);
      if (OptimizationMode.balanced.equals(optimizationMode)) {

        List<WebPageContent> rankedWebPageContents = vectorRankerService.filter(webPageContents, quesiton, 1);
        rankedWebPageContents = jinaReaderService.spider(webPageContents);
        webPageContents.set(0, rankedWebPageContents.get(0));

      } else if (OptimizationMode.quality.equals(optimizationMode)) {
        // 质量模式下先过滤，再异步补全页面内容
        webPageContents = aiRankerService.filter(webPageContents, quesiton, 6);
        webPageContents = jinaReaderService.spiderAsync(webPageContents);
      }

      chatParamVo.setSources(webPageContents);
      // 将搜索结果转换为 JSON 格式保存到数据库中（便于记录历史消息）
      PGobject pgObject = PgObjectUtils.json(webPageContents);
      Db.updateBySql("update max_search_chat_message set sources=? where id=?", pgObject, questionMessageId);

      List<WebPageSource> sources = new ArrayList<>();

      for (WebPageContent webPageConteont : webPageContents) {
        sources.add(new WebPageSource(webPageConteont.getTitle(), webPageConteont.getUrl(), webPageConteont.getContent()));
      }

      String host = channelContext.getString(RequestHeaderKey.Host);
      if (host == null) {
        host = "//127.0.0.1";
      } else {
        host = "//" + host;
      }
      sources.add(new WebPageSource("All Sources", host + "/sources/" + questionMessageId));
      // 返回 sources 数据给客户端
      ChatWsRespVo<List<WebPageSource>> chatRespVo = new ChatWsRespVo<>();
      chatRespVo.setType("sources").setData(sources).setMessageId(answerMessageId);

      // 通过 WebSocket or sse 返回搜索结果引用信息给客户端
      if (channelContext != null) {
        byte[] jsonBytes = FastJson2Utils.toJSONBytes(chatRespVo);
        if (reqMessageVo.isSse()) {
          Tio.bSend(channelContext, new SsePacket(jsonBytes));
        } else {
          Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
        }
      }

      // 拼接所有搜索结果内容，用于生成提示词
      StringBuffer markdown = new StringBuffer();
      for (int i = 0; i < webPageContents.size(); i++) {
        WebPageContent webPageContent = webPageContents.get(i);
        String sourceContent = webPageContent.getContent();
        if (StrUtil.isBlank(sourceContent)) {
          sourceContent = webPageContent.getDescription();
        }
        String sourceFormat = "source %d %s %s  ";
        markdown.append(String.format(sourceFormat, (i + 1), webPageContent.getUrl(), sourceContent));
      }

      // 使用模板引擎生成提示词，提示词中包含当前日期和搜索结果上下文
      String isoTimeStr = DateTimeFormatter.ISO_INSTANT.format(Instant.now());
      Kv kv = Kv.by("date", isoTimeStr).set("context", markdown);
      inputPrompt = PromptEngine.renderToString("WebSearchResponsePrompt.txt", kv);
      log.info("deepkseek:{}", inputPrompt);
    }
    chatParamVo.setSystemPrompt(inputPrompt);
    return predictService.predict(channelContext, reqMessageVo, chatParamVo);
  }
}

说明

  • 判断搜索条件：当 copilotEnabled 为 true 时，会优先使用重写后的问题（若有）进行搜索。
  • 调用搜索 API：通过 maxSearchSearchService.search 获取相关网页内容。
  • 优化模式处理：根据用户设置的优化模式（balanced 或 quality）分别进行搜索结果的过滤和页面内容补全。
  • 返回引用信息：构造包含各个数据源信息的列表，并以 sources 类型的消息返回给前端。
  • 生成提示词：将所有搜索结果内容拼接后，通过模板引擎生成最终提示词，供后续回答生成使用。

2.4 PredictService

该服务用于封装预测服务的调用，根据具体需求选择使用 GeminiPredictService 或 DeepSeekPredictService，目前默认调用 DeepSeekPredictService。

package com.litongjava.perplexica.services;

import com.litongjava.jfinal.aop.Aop;
import com.litongjava.perplexica.vo.ChatParamVo;
import com.litongjava.perplexica.vo.ChatWsReqMessageVo;
import com.litongjava.tio.core.ChannelContext;

import okhttp3.Call;

public class PredictService {
  private GeminiPredictService geminiPredictService = Aop.get(GeminiPredictService.class);
  private DeepSeekPredictService deepSeekPredictService = Aop.get(DeepSeekPredictService.class);

  public Call predict(ChannelContext channelContext, ChatWsReqMessageVo reqMessageVo, ChatParamVo chatParamVo) {
    return deepSeekPredictService.predict(channelContext, reqMessageVo, chatParamVo);
  }

}

说明

  • 此服务起到简单的封装作用，目前直接调用 DeepSeekPredictService 来生成回答。

2.5 DeepSeekPredictService

该服务基于 DeepSeek 模型生成回答，构造请求后以流式方式调用 OpenAiClient，并通过 DeepSeekSseCallback 处理返回的 SSE 流数据。

package com.litongjava.perplexica.services;

import java.util.ArrayList;
import java.util.List;

import com.litongjava.openai.chat.OpenAiChatMessage;
import com.litongjava.openai.chat.OpenAiChatRequestVo;
import com.litongjava.openai.client.OpenAiClient;
import com.litongjava.perplexica.callback.DeepSeekSseCallback;
import com.litongjava.perplexica.vo.ChatParamVo;
import com.litongjava.perplexica.vo.ChatWsReqMessageVo;
import com.litongjava.perplexica.vo.ChatWsRespVo;
import com.litongjava.tio.core.ChannelContext;
import com.litongjava.tio.core.Tio;
import com.litongjava.tio.http.common.sse.SsePacket;
import com.litongjava.tio.utils.SystemTimer;
import com.litongjava.tio.utils.environment.EnvUtils;
import com.litongjava.tio.utils.json.FastJson2Utils;
import com.litongjava.tio.websocket.common.WebSocketResponse;
import com.litongjava.volcengine.VolcEngineConst;
import com.litongjava.volcengine.VolcEngineModels;

import lombok.extern.slf4j.Slf4j;
import okhttp3.Call;
import okhttp3.Callback;

@Slf4j
public class DeepSeekPredictService {

  public Call predict(ChannelContext channelContext, ChatWsReqMessageVo reqMessageVo, ChatParamVo chatParamVo) {
    String systemPrompt = chatParamVo.getSystemPrompt();
    Long sessionId = reqMessageVo.getMessage().getChatId();
    Long questionMessageId = reqMessageVo.getMessage().getMessageId();
    String content = reqMessageVo.getMessage().getContent();
    Long answerMessageId = chatParamVo.getAnswerMessageId();

    List<OpenAiChatMessage> contents = new ArrayList<>();
    if (systemPrompt != null) {
      contents.add(new OpenAiChatMessage("system", systemPrompt));
      log.info("deepkseek:{}", systemPrompt);
    }

    List<List<String>> history = reqMessageVo.getHistory();
    if (history != null && history.size() > 0) {
      for (int i = 0; i < history.size(); i++) {
        String role = history.get(i).get(0);
        String message = history.get(i).get(1);
        if ("human".equals(role)) {
          role = "user";
        } else {
          role = "assistant";
        }
        contents.add(new OpenAiChatMessage(role, message));
      }
    }

    contents.add(new OpenAiChatMessage("user", content));

    OpenAiChatRequestVo chatRequestVo = new OpenAiChatRequestVo().setModel(VolcEngineModels.DEEPSEEK_R1_250120)
        //
        .setMessages(contents);
    chatRequestVo.setStream(true);

    // 5. 向前端通知一个空消息，标识搜索结束，开始推理
    //{"type":"message","data":"", "messageId": "32fcbbf251337c"}
    ChatWsRespVo<String> chatVo = ChatWsRespVo.message(answerMessageId, "");
    byte[] jsonBytes = FastJson2Utils.toJSONBytes(chatVo);
    if (channelContext != null) {
      if (reqMessageVo.isSse()) {
        Tio.bSend(channelContext, new SsePacket(jsonBytes));
      } else {
        Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
      }
    }

    long start = SystemTimer.currTime;
    Callback callback = new DeepSeekSseCallback(channelContext, reqMessageVo, chatParamVo, start);
    String apiKey = EnvUtils.getStr("VOLCENGINE_API_KEY");
    Call call = OpenAiClient.chatCompletions(VolcEngineConst.BASE_URL, apiKey, chatRequestVo, callback);
    return call;
  }

}

说明

  • 构造消息列表：根据系统提示、历史对话和当前用户输入构造完整的对话消息列表。
  • 生成请求：使用 VolcEngine 提供的 DeepSeek 模型（DEEPSEEK_R1_250120）构造 OpenAiChatRequestVo，并启用流式传输。
  • 通知前端：发送一个空消息通知前端搜索阶段结束，开始生成回答。
  • 调用预测接口：利用 OpenAiClient 的 chatCompletions 方法发起请求，并将回调设置为 DeepSeekSseCallback。

2.6 DeepSeekSseCallback

该回调类负责处理 DeepSeek 模型返回的 SSE 流数据，将解析后的数据实时发送给客户端，同时处理错误、进度、和结束信号，并保存最终生成的回答。

package com.litongjava.perplexica.callback;

import java.io.IOException;
import java.util.List;

import com.jfinal.kit.Kv;
import com.litongjava.openai.chat.ChatResponseDelta;
import com.litongjava.openai.chat.Choice;
import com.litongjava.openai.chat.OpenAiChatResponseVo;
import com.litongjava.perplexica.can.ChatWsStreamCallCan;
import com.litongjava.perplexica.model.MaxSearchChatMessage;
import com.litongjava.perplexica.vo.ChatParamVo;
import com.litongjava.perplexica.vo.ChatWsReqMessageVo;
import com.litongjava.perplexica.vo.ChatWsRespVo;
import com.litongjava.tio.core.ChannelContext;
import com.litongjava.tio.core.Tio;
import com.litongjava.tio.http.common.sse.SsePacket;
import com.litongjava.tio.http.server.util.SseEmitter;
import com.litongjava.tio.utils.json.FastJson2Utils;
import com.litongjava.tio.websocket.common.WebSocketResponse;

import lombok.extern.slf4j.Slf4j;
import okhttp3.Call;
import okhttp3.Callback;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okio.BufferedSource;

@Slf4j
public class DeepSeekSseCallback implements Callback {
  private ChannelContext channelContext;
  private ChatWsReqMessageVo reqVo;
  private ChatParamVo chatParamVo;
  private long start;
  private Long sessionId;
  private long answerMessageId;

  public DeepSeekSseCallback(ChannelContext channelContext, ChatWsReqMessageVo reqMessageVo, ChatParamVo chatParamVo, long start) {
    this.channelContext = channelContext;
    this.reqVo = reqMessageVo;
    this.chatParamVo = chatParamVo;
    this.start = start;
    Long sessionId = reqVo.getMessage().getChatId();
    this.sessionId = sessionId;
    this.answerMessageId = chatParamVo.getAnswerMessageId();
  }

  @Override
  public void onFailure(Call call, IOException e) {
    ChatWsRespVo<String> error = ChatWsRespVo.error("CHAT_ERROR", e.getMessage());
    byte[] jsonBytes = FastJson2Utils.toJSONBytes(error);
    if (reqVo.isSse()) {
      Tio.bSend(channelContext, new SsePacket(jsonBytes));
      ChatWsStreamCallCan.remove(sessionId + "");
      SseEmitter.closeSeeConnection(channelContext);

    } else {
      WebSocketResponse packet = new WebSocketResponse(jsonBytes);
      Tio.bSend(channelContext, packet);
    }

  }

  @Override
  public void onResponse(Call call, Response response) throws IOException {
    if (!response.isSuccessful()) {
      String string = response.body().string();
      String message = "Chat model response an unsuccessful message:" + string;
      log.error("message:{}", message);
      ChatWsRespVo<String> data = ChatWsRespVo.error("STREAM_ERROR", message);
      byte[] jsonBytes = FastJson2Utils.toJSONBytes(data);
      if (reqVo.isSse()) {
        Tio.bSend(channelContext, new SsePacket(jsonBytes));
      } else {
        Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
      }
      return;
    }

    try (ResponseBody responseBody = response.body()) {
      if (responseBody == null) {
        String message = "response body is null";
        log.error(message);
        ChatWsRespVo<String> data = ChatWsRespVo.progress(message);
        byte[] jsonBytes = FastJson2Utils.toJSONBytes(data);
        if (reqVo.isSse()) {
          Tio.bSend(channelContext, new SsePacket(jsonBytes));
        } else {
          Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
        }

        return;
      }
      StringBuffer completionContent = onResponseSuccess(channelContext, answerMessageId, start, responseBody);
      // save user mesasge
      new MaxSearchChatMessage().setId(answerMessageId).setChatId(sessionId)
          //
          .setRole("assistant").setContent(completionContent.toString())
          //
          .save();

      Kv end = Kv.by("type", "messageEnd").set("messageId", answerMessageId);
      byte[] jsonBytes = FastJson2Utils.toJSONBytes(end);
      if (reqVo.isSse()) {
        Tio.bSend(channelContext, new SsePacket(jsonBytes));
      } else {
        Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
      }

      // 关闭连接
      long endTime = System.currentTimeMillis();
      log.info("finish llm in {} (ms)", (endTime - start));

      //log.info("completionContent:{}", completionContent);
      if (completionContent != null && !completionContent.toString().isEmpty()) {
        //函数调用处理
      }
    }
    ChatWsStreamCallCan.remove(sessionId + "");
    if (reqVo.isSse()) {
      // 手动移除连接
      SseEmitter.closeSeeConnection(channelContext);
    }
  }

  /**
   * 处理ChatGPT成功响应
   *
   * @param channelContext 通道上下文
   * @param responseBody    响应体
   * @return 完整内容
   * @throws IOException
   */
  public StringBuffer onResponseSuccess(ChannelContext channelContext, Long answerMessageId, Long start, ResponseBody responseBody) throws IOException {
    StringBuffer completionContent = new StringBuffer();
    BufferedSource source = responseBody.source();
    String line;

    while ((line = source.readUtf8Line()) != null) {
      if (line.length() < 1) {
        continue;
      }
      // 处理数据行
      if (line.length() > 6) {
        String data = line.substring(6);
        if (data.endsWith("}")) {
          OpenAiChatResponseVo chatResponse = FastJson2Utils.parse(data, OpenAiChatResponseVo.class);
          List<Choice> choices = chatResponse.getChoices();
          if (!choices.isEmpty()) {
            ChatResponseDelta delta = choices.get(0).getDelta();

            String reasoning_content = delta.getReasoning_content();
            if (reasoning_content != null && !reasoning_content.isEmpty()) {
              ChatWsRespVo<String> vo = ChatWsRespVo.reasoning(answerMessageId, reasoning_content);
              byte[] jsonBytes = FastJson2Utils.toJSONBytes(vo);
              if (reqVo.isSse()) {
                Tio.bSend(channelContext, new SsePacket(jsonBytes));
              } else {
                Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
              }
            }


            String part = delta.getContent();
            if (part != null && !part.isEmpty()) {
              completionContent.append(part);
              ChatWsRespVo<String> vo = ChatWsRespVo.message(answerMessageId, part);
              byte[] jsonBytes = FastJson2Utils.toJSONBytes(vo);
              if (reqVo.isSse()) {
                Tio.bSend(channelContext, new SsePacket(jsonBytes));
              } else {
                Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
              }
            }

          }
        } else if (": keep-alive".equals(line)) {
          ChatWsRespVo<String> vo = ChatWsRespVo.keepAlive(answerMessageId);
          byte[] jsonBytes = FastJson2Utils.toJSONBytes(vo);
          if (reqVo.isSse()) {

          } else {
            Tio.bSend(channelContext, new WebSocketResponse(jsonBytes));
          }

        } else {
          log.info("Data does not end with }:{}", line);
          //{"type":"messageEnd","messageId":"654b8bdb25e853"}

        }
      }
    }
    return completionContent;
  }
}

说明

  • onFailure 方法：请求失败时，发送错误消息给前端，并移除对应的流调用，关闭 SSE 连接。
  • onResponse 方法：
    • 检查响应状态码，若不成功则返回错误消息。
    • 调用 onResponseSuccess 逐行解析响应数据，提取 JSON 中的 delta 信息，对应内容通过 SSE 分为 reasoning 与 message 类型实时发送。
    • 最后发送 messageEnd 消息，并将生成的回答保存到数据库。
  • onResponseSuccess 方法：负责解析流数据，提取每一行数据中的内容，并根据内容类型实时发送给前端。

3. 总结

本文档详细介绍了 ai-search 搜索代码的 SSE 版本实现，包括：

  • 接口说明：定义了 SSE 请求、响应头及消息流中各类型消息的格式与作用。
  • 核心代码实现：涵盖了请求处理、业务逻辑分发、搜索结果获取与优化、以及 DeepSeek 模型的调用与流式响应处理。

通过以上代码和详尽的解释，读者可以了解如何实现基于 SSE 的实时数据传输、如何整合搜索和 AI 模型进行回答生成，以及各模块间的调用关系和数据流转。希望这篇文档能够帮助您在实际项目中借鉴和改进类似的技术方案。