item3 - 11 聊天模块

import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.TextMessage;

@Component
public class MyHandler extends TextWebSocketHandler {

  private final ConcurrentHashMap<String, WebSocketSession> sessions = new ConcurrentHashMap<>();

  // 1. 连接建立后触发
  @Override
  public void afterConnectionEstablished(WebSocketSession session) {
    System.out.println("用户来了，ID: " + session.getId());
    // 通常在这里把 session 存到一个 Map 里，方便后续推送
    sessions.put(session.getId(), session);
  }

  // 2. 收到消息时触发 (最重要)
  @Override
  protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
    String payload = message.getPayload(); // 获取用户发的内容
    System.out.println("收到消息: " + payload);

    // 回复消息 (这就是原生写法，直接 send)
    session.sendMessage(new TextMessage("服务器收到了: " + payload));
  }

  // 3. 连接关闭后触发
  @Override
  public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
    System.out.println("用户走了");
    // 在这里把 session 从 Map 里移除
    sessions.remove(session.getId());
  }

  // 服务端主动推送/广播
  public void broadcast(String text) {
    TextMessage msg = new TextMessage(text);
    sessions.values().forEach(s -> {
        try {
            if (s.isOpen()) s.sendMessage(msg);
        } catch (IOException ignored) {}
    });
  }
}

import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.beans.factory.annotation.Autowired;

@Configuration
@EnableWebSocket // 关键注解：开启 WebSocket 功能
public class WebSocketConfig implements WebSocketConfigurer {
  @Autowired
  private MyHandler myHandler; // 注入第一步写的 Handler

  @Override
  public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
      // 核心配置：
      // 1. 绑定 handler
      // 2. 设定路径 "/ws/chat"
      // 3. 允许跨域 (setAllowedOrigins)
      registry.addHandler(myHandler, "/ws/chat")
              .setAllowedOrigins("*"); 
  }
}

import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;

@Configuration
@EnableWebSocketMessageBroker // 开启 STOMP 代理
public class WebSocketStompConfig implements WebSocketMessageBrokerConfigurer {

  @Override
  public void registerStompEndpoints(StompEndpointRegistry registry) {
      // 1. 设定连接端点，前端连接这个地址: ws://localhost:8080/ws-stomp
      registry.addEndpoint("/ws-stomp")
              .setAllowedOriginPatterns("*") // 允许跨域
              .withSockJS(); // 开启 SockJS 支持（可选，推荐）
  }

  @Override
  public void configureMessageBroker(MessageBrokerRegistry registry) {
      // 2. 配置消息代理
      // 凡是客户端订阅 "/topic" 或 "/queue" 开头的，都由代理处理（发回给客户端）
      registry.enableSimpleBroker("/topic", "/queue");
      
      // 凡是客户端发往 "/app" 开头的消息，都会路由到 @Controller 里的方法
      registry.setApplicationDestinationPrefixes("/app");
  }
}

import org.springframework.messaging.handler.annotation.MessageMapping;
import org.springframework.messaging.handler.annotation.SendTo;
import org.springframework.stereotype.Controller;

@Controller
public class StompChatController {

  // 1. 接收消息
  // 对应前端发送地址: /app/hello (因为配置了前缀 /app)
  @MessageMapping("/hello")
  // 2. 广播消息
  // 处理完后，返回值会自动发给订阅了 /topic/greetings 的所有人
  @SendTo("/topic/greetings")
  public String greeting(String message) {
      System.out.println("收到前端消息: " + message);
      return "服务器回复: " + message;
  }
}

private final WebClient webClient;
private final String apiKey;
private final String model;
private final AiProperties aiProperties;
private static final Logger logger = LoggerFactory.getLogger(DeepSeekClient.class);

public DeepSeekClient(@Value("${deepseek.api.url}") String apiUrl,
  @Value("${deepseek.api.key}") String apiKey,
  @Value("${deepseek.api.model}") String model,
  AiProperties aiProperties) {
  WebClient.Builder builder = WebClient.builder().baseUrl(apiUrl);

  // 只有当 API key 不为空时才添加 Authorization header
  if (apiKey != null && !apiKey.trim().isEmpty()) {
    builder.defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + apiKey);
  }

  this.webClient = builder.build();
  this.apiKey = apiKey;
  this.model = model;
  this.aiProperties = aiProperties;
}

public void streamResponse(String userMessage, String context, List<Map<String, String>> history, Consumer<String> onChunk Consumer<Throwable> onError) {
  
  Map<String, Object> request = buildRequest(userMessage, context, history);
  
  webClient.post()
    .uri("/chat/completions")
    .contentType(MediaType.APPLICATION_JSON)
    .bodyValue(request)
    .retrieve()
    .bodyToFlux(String.class)
    .subscribe(
        chunk -> processChunk(chunk, onChunk),
        onError
    );
}

webClient.post()
  .uri("/chat/completions")
  .contentType(MediaType.APPLICATION_JSON)
  .bodyValue(request)
  .retrieve()

.bodyToFlux(String.class)

.subscribe(
  chunk -> processChunk(chunk, onChunk),
  onError
);

private Map<String, Object> buildRequest(String userMessage, String context, List<Map<String, String>> history) {
  logger.info("构建请求，用户消息：{}，上下文长度：{}，历史消息数：{}", 
    userMessage, 
    context != null ? context.length() : 0, 
    history != null ? history.size() : 0);
  
  Map<String, Object> request = new java.util.HashMap<>();
  request.put("model", model);
  request.put("messages", buildMessages(userMessage, context, history));
  request.put("stream", true);
  // 生成参数
  AiProperties.Generation gen = aiProperties.getGeneration();
  if (gen.getTemperature() != null) {
      request.put("temperature", gen.getTemperature());
  }
  if (gen.getTopP() != null) {
      request.put("top_p", gen.getTopP());
  }
  if (gen.getMaxTokens() != null) {
      request.put("max_tokens", gen.getMaxTokens());
  }
  return request;
}

private List<Map<String, String>> buildMessages(String userMessage, String context, List<Map<String, String>> history) {
  List<Map<String, String>> messages = new ArrayList<>();

  AiProperties.Prompt promptCfg = aiProperties.getPrompt();

  // 1. 构建统一的 system 指令（规则 + 参考信息）
  StringBuilder sysBuilder = new StringBuilder();
  String rules = promptCfg.getRules();
  if (rules != null) {
      sysBuilder.append(rules).append("\n\n");
  }

  String refStart = promptCfg.getRefStart() != null ? promptCfg.getRefStart() : "<<REF>>";
  String refEnd = promptCfg.getRefEnd() != null ? promptCfg.getRefEnd() : "<<END>>";
  sysBuilder.append(refStart).append("\n");

  if (context != null && !context.isEmpty()) {
      sysBuilder.append(context);
  } else {
      String noResult = promptCfg.getNoResultText() != null ? promptCfg.getNoResultText() : "（本轮无检索结果）";
      sysBuilder.append(noResult).append("\n");
  }

  sysBuilder.append(refEnd);

  String systemContent = sysBuilder.toString();
  messages.add(Map.of(
      "role", "system",
      "content", systemContent
  ));
  logger.debug("添加了系统消息，长度: {}", systemContent.length());

  // 2. 追加历史消息（若有）
  if (history != null && !history.isEmpty()) {
      messages.addAll(history);
  }

  // 3. 当前用户问题
  messages.add(Map.of(
      "role", "user",
      "content", userMessage
  ));

  return messages;
}

private void processChunk(String chunk, Consumer<String> onChunk) {
  try {
      // 检查是否是结束标记
      if ("[DONE]".equals(chunk)) {
          logger.debug("对话结束");
          return;
      }
      
      // 直接解析 JSON
      ObjectMapper mapper = new ObjectMapper();
      JsonNode node = mapper.readTree(chunk);
      String content = node.path("choices")
                          .path(0)
                          .path("delta")
                          .path("content")
                          .asText("");
      
      if (!content.isEmpty()) {
          onChunk.accept(content);
      }
  } catch (Exception e) {
      logger.error("处理数据块时出错: {}", e.getMessage(), e);
  }
}