Commit 875cfa10c6fbbd4ab224fe87b973daee84b7d7e1

Authored by YunaiV
1 parent 4ee638db

【代码新增】AI:适配 提供 OllamaChatModelTests,用于本地 ollma 模型的调用

yudao-module-ai/yudao-spring-boot-starter-ai/src/test/java/cn/iocoder/yudao/framework/ai/chat/OllamaChatModelTests.java 0 → 100644
  1 +package cn.iocoder.yudao.framework.ai.chat;
  2 +
  3 +import org.junit.jupiter.api.Disabled;
  4 +import org.junit.jupiter.api.Test;
  5 +import org.springframework.ai.chat.messages.Message;
  6 +import org.springframework.ai.chat.messages.SystemMessage;
  7 +import org.springframework.ai.chat.messages.UserMessage;
  8 +import org.springframework.ai.chat.model.ChatResponse;
  9 +import org.springframework.ai.chat.prompt.Prompt;
  10 +import org.springframework.ai.ollama.OllamaChatModel;
  11 +import org.springframework.ai.ollama.api.OllamaApi;
  12 +import org.springframework.ai.ollama.api.OllamaOptions;
  13 +import reactor.core.publisher.Flux;
  14 +
  15 +import java.util.ArrayList;
  16 +import java.util.List;
  17 +
  18 +/**
  19 + * {@link OllamaChatModel} 集成测试
  20 + *
  21 + * @author 芋道源码
  22 + */
  23 +public class OllamaChatModelTests {
  24 +
  25 + private final OllamaChatModel chatModel = OllamaChatModel.builder()
  26 + .ollamaApi(new OllamaApi("http://127.0.0.1:11434")) // Ollama 服务地址
  27 + .defaultOptions(OllamaOptions.builder()
  28 +// .model("qwen") // 模型(https://ollama.com/library/qwen)
  29 + .model("deepseek-r1") // 模型(https://ollama.com/library/deepseek-r1)
  30 + .build())
  31 + .build();
  32 +
  33 + @Test
  34 + @Disabled
  35 + public void testCall() {
  36 + // 准备参数
  37 + List<Message> messages = new ArrayList<>();
  38 + messages.add(new SystemMessage("你是一个优质的文言文作者,用文言文描述着各城市的人文风景。"));
  39 + messages.add(new UserMessage("1 + 1 = ?"));
  40 +
  41 + // 调用
  42 + ChatResponse response = chatModel.call(new Prompt(messages));
  43 + // 打印结果
  44 + System.out.println(response);
  45 + System.out.println(response.getResult().getOutput());
  46 + }
  47 +
  48 + @Test
  49 + @Disabled
  50 + public void testStream() {
  51 + // 准备参数
  52 + List<Message> messages = new ArrayList<>();
  53 + messages.add(new SystemMessage("你是一个优质的文言文作者,用文言文描述着各城市的人文风景。"));
  54 + messages.add(new UserMessage("1 + 1 = ?"));
  55 +
  56 + // 调用
  57 + Flux<ChatResponse> flux = chatModel.stream(new Prompt(messages));
  58 + // 打印结果
  59 + flux.doOnNext(response -> {
  60 +// System.out.println(response);
  61 + System.out.println(response.getResult().getOutput());
  62 + }).then().block();
  63 + }
  64 +
  65 +}