Spring AI + ollama 本地搭建聊天 AI

Spring AI + ollama 本地搭建聊天 AI

不知道怎么搭建 ollama 的可以查看上一篇Spring AI 初学。

项目可以查看gitee

前期准备

添加依赖

创建 SpringBoot 项目，添加主要相关依赖（spring-boot-starter-web、spring-ai-ollama-spring-boot-starter）

Spring AI supports Spring Boot 3.2.x and 3.3.x

Spring Boot 3.2.11 requires at least Java 17 and is compatible with versions up to and including Java 23

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-web</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-ollama-spring-boot-starter</artifactId>
    <version>1.0.0-M3</version>
</dependency>

配置文件

application.properties、yml配置文件中添加，也可以在项目中指定模型等参数，具体参数可以参考 OllamaChatProperties

# properties，模型 qwen2.5:14b 根据自己下载的模型而定
spring.ai.ollama.chat.options.model=qwen2.5:14b

#yml
spring:
  ai:
    ollama:
      chat:
        model: qwen2.5:14b

聊天实现

主要使用 org.springframework.ai.chat.memory.ChatMemory 接口保存对话信息。

一、采用 Java 缓存对话信息

支持功能：聊天对话、切换对话、删除对话

controller

import com.yb.chatai.domain.ChatParam;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.web.bind.annotation.*;

import java.util.UUID;

/*
 *@title Controller
 *@description 使用内存进行对话
 *@author yb
 *@version 1.0
 *@create 2024/11/12 14:39
 */
@Controller
public class ChatController {

    //注入模型，配置文件中的模型，或者可以在方法中指定模型
    @Resource
    private OllamaChatModel model;

    //聊天 client
    private ChatClient chatClient;

    // 模拟数据库存储会话和消息
    private final ChatMemory chatMemory = new InMemoryChatMemory();

    //首页
    @GetMapping("/index")
    public String index(){
        return "index";
    }

    //开始聊天，生成唯一 sessionId
    @GetMapping("/start")
    public String start(Model model){
        //新建聊天模型
//        OllamaOptions options = OllamaOptions.builder();
//        options.setModel("qwen2.5:14b");
//        OllamaChatModel chatModel = new OllamaChatModel(new OllamaApi(), options);
        //创建随机会话 ID
        String sessionId = UUID.randomUUID().toString();
        model.addAttribute("sessionId", sessionId);
        //创建聊天client
        chatClient = ChatClient.builder(this.model).defaultAdvisors(new MessageChatMemoryAdvisor(chatMemory, sessionId, 10)).build();
        return "chatPage";
    }

    //聊天
    @PostMapping("/chat")
    @ResponseBody
    public String chat(@RequestBody ChatParam param){
        //直接返回
        return chatClient.prompt(param.getUserMsg()).call().content();
    }

    //删除聊天
    @DeleteMapping("/clear/{id}")
    @ResponseBody
    public void clear(@PathVariable("id") String sessionId){
        chatMemory.clear(sessionId);
    }

}

效果图

二、采用数据库保存对话信息

支持功能：聊天对话、切换对话、删除对话、撤回消息

实体类

import lombok.Data;

import java.util.Date;

@Data
public class ChatEntity {

    private String id;

    /** 会话id */
    private String sessionId;

    /** 会话内容 */
    private String content;

    /** AI、人 */
    private String type;

    /** 创建时间 */
    private Date time;

    /** 是否删除，Y-是 */
    private String beDeleted;

    /** AI会话时，获取人对话ID */
    private String userChatId;

}

configuration

import com.yb.chatai.domain.ChatEntity;
import com.yb.chatai.service.IChatService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.context.annotation.Configuration;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/*
 *@title DBMemory
 *@description 实现 ChatMemory，注入 spring，方便采用 service 方法
 *@author yb
 *@version 1.0
 *@create 2024/11/12 16:15
 */
@Configuration
public class DBMemory implements ChatMemory {

    @Resource
    private IChatService chatService;

    @Override
    public void add(String conversationId, List<Message> messages) {
        for (Message message : messages) {
            chatService.saveMessage(conversationId, message.getContent(), message.getMessageType().getValue());
        }
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<ChatEntity> list = chatService.getLastN(conversationId, lastN);
        if(list != null && !list.isEmpty()) {
            return list.stream().map(l -> {
                Message message = null;
                if (MessageType.ASSISTANT.getValue().equals(l.getType())) {
                    message = new AssistantMessage(l.getContent());
                } else if (MessageType.USER.getValue().equals(l.getType())) {
                    message = new UserMessage(l.getContent());
                }
                return message;
            }).collect(Collectors.<Message>toList());
        }else {
            return new ArrayList<>();
        }
    }

    @Override
    public void clear(String conversationId) {
        chatService.clear(conversationId);
    }
}

services实现类

import com.yb.chatai.domain.ChatEntity;
import com.yb.chatai.service.IChatService;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.stereotype.Service;

import java.util.*;

/*
 *@title ChatServiceImpl
 *@description 保存用户会话 service 实现类
 *@author yb
 *@version 1.0
 *@create 2024/11/12 15:50
 */
@Service
public class ChatServiceImpl implements IChatService {

    Map<String, List<ChatEntity>> map = new HashMap<>();

    @Override
    public void saveMessage(String sessionId, String content, String type) {
        ChatEntity entity = new ChatEntity();
        entity.setId(UUID.randomUUID().toString());
        entity.setContent(content);
        entity.setSessionId(sessionId);
        entity.setType(type);
        entity.setTime(new Date());
        //改成常量
        entity.setBeDeleted("N");
        if(MessageType.ASSISTANT.getValue().equals(type)){
            entity.setUserChatId(getLastN(sessionId, 1).get(0).getId());
        }
        //todo 保存数据库
        //模拟保存到数据库
        List<ChatEntity> list = map.getOrDefault(sessionId, new ArrayList<>());
        list.add(entity);
        map.put(sessionId, list);
    }

    @Override
    public List<ChatEntity> getLastN(String sessionId, Integer lastN) {
        //todo 从数据库获取
        //模拟从数据库获取
        List<ChatEntity> list = map.get(sessionId);
        return list != null ? list.stream().skip(Math.max(0, list.size() - lastN)).toList() : List.of();
    }

    @Override
    public void clear(String sessionId) {
        //todo 数据库更新 beDeleted 字段
        map.put(sessionId, new ArrayList<>());
    }

    @Override
    public void deleteById(String id) {
        //todo 数据库直接将该 id 数据 beDeleted 改成 Y
        for (Map.Entry<String, List<ChatEntity>> next : map.entrySet()) {
            List<ChatEntity> list = next.getValue();
            list.removeIf(chat -> id.equals(chat.getId()) || id.equals(chat.getUserChatId()));
        }
    }
}

controller

import com.yb.chatai.configuration.DBMemory;
import com.yb.chatai.domain.ChatEntity;
import com.yb.chatai.domain.ChatParam;
import com.yb.chatai.service.IChatService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.web.bind.annotation.*;

import java.util.List;
import java.util.UUID;

/*
 *@title ChatController2
 *@description 使用数据库（缓存）进行对话
 *@author yb
 *@version 1.0
 *@create 2024/11/12 16:12
 */
@Controller
public class ChatController2 {

    //注入模型，配置文件中的模型，或者可以在方法中指定模型
    @Resource
    private OllamaChatModel model;

    //聊天 client
    private ChatClient chatClient;

    //操作聊天信息service
    @Resource
    private IChatService chatService;

    //会话存储方式
    @Resource
    private DBMemory dbMemory;

    //开始聊天，生成唯一 sessionId
    @GetMapping("/start2")
    public String start(Model model){
        //新建聊天模型
//        OllamaOptions options = OllamaOptions.builder();
//        options.setModel("qwen2.5:14b");
//        OllamaChatModel chatModel = new OllamaChatModel(new OllamaApi(), options);
        //创建随机会话 ID
        String sessionId = UUID.randomUUID().toString();
        model.addAttribute("sessionId", sessionId);
        //创建聊天 client
        chatClient = ChatClient.builder(this.model).defaultAdvisors(new MessageChatMemoryAdvisor(dbMemory, sessionId, 10)).build();
        return "chatPage2";
    }

    //切换会话，需要传入 sessionId
    @GetMapping("/exchange2/{id}")
    public String exchange(@PathVariable("id")String sessionId){
        //切换聊天 client
        chatClient = ChatClient.builder(this.model).defaultAdvisors(new MessageChatMemoryAdvisor(dbMemory, sessionId, 10)).build();
        return "chatPage2";
    }

    //聊天
    @PostMapping("/chat2")
    @ResponseBody
    public List<ChatEntity> chat(@RequestBody ChatParam param){
        //todo 判断 AI 是否返回会话，从而判断用户是否可以输入
        chatClient.prompt(param.getUserMsg()).call().content();
        //获取返回最新两条，一条用户问题（用户获取用户发送ID），一条 AI 返回结果
        return chatService.getLastN(param.getSessionId(), 2);
    }

    //撤回消息
    @DeleteMapping("/revoke2/{id}")
    @ResponseBody
    public void revoke(@PathVariable("id") String id){
        chatService.deleteById(id);
    }

    //清空消息
    @DeleteMapping("/del2/{id}")
    @ResponseBody
    public void clear(@PathVariable("id") String sessionId){
        dbMemory.clear(sessionId);
    }

}

效果图

总结

主要实现 org.springframework.ai.chat.memory.ChatMemory 方法，实际项目过程需要实现该接口重写方法。