BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 0139317ac4ee3a05be5ab9d432191460c1107ae3

Authored by YunaiV
1 parent 1feff2b1

【代码优化】AI：完善 LlamaChatModelTests 单测，方便大家快速体验

Inline Side-by-side
Showing 2 changed files with 66 additions and 3 deletions
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/core/factory/AiModelFactoryImpl.java
  View file @0139317
... ... @@ -55,8 +55,6 @@ public class AiModelFactoryImpl implements AiModelFactory {
55 55 return Singleton.get(cacheKey, (Func0<ChatModel>) () -> {
56 56 //noinspection EnhancedSwitchMigration
57 57 switch (platform) {
58   - case OLLAMA:
59   - return buildOllamaChatClient(url);
60 58 case YI_YAN:
61 59 return buildYiYanChatClient(apiKey);
62 60 case XING_HUO:
... ... @@ -67,6 +65,8 @@ public class AiModelFactoryImpl implements AiModelFactory {
67 65 return buildDeepSeekChatClient(apiKey);
68 66 case OPENAI:
69 67 return buildOpenAiChatModel(apiKey, url);
  68 + case OLLAMA:
  69 + return buildOllamaChatModel(url);
70 70 default:
71 71 throw new IllegalArgumentException(StrUtil.format("未知平台({})", platform));
72 72 }
... ... @@ -163,7 +163,7 @@ public class AiModelFactoryImpl implements AiModelFactory {
163 163 /**
164 164 * 可参考 {@link OllamaAutoConfiguration}
165 165 */
166   - private static OllamaChatModel buildOllamaChatClient(String url) {
  166 + private static OllamaChatModel buildOllamaChatModel(String url) {
167 167 OllamaApi ollamaApi = new OllamaApi(url);
168 168 return new OllamaChatModel(ollamaApi);
169 169 }
... ...
yudao-module-ai/yudao-spring-boot-starter-ai/src/test/java/cn/iocoder/yudao/framework/ai/chat/LlamaChatModelTests.java 0 → 100644
  View file @0139317
  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.OllamaModel;
  13 +import org.springframework.ai.ollama.api.OllamaOptions;
  14 +import reactor.core.publisher.Flux;
  15 +
  16 +import java.util.ArrayList;
  17 +import java.util.List;
  18 +
  19 +/**
  20 + * {@link OllamaChatModel} 集成测试
  21 + *
  22 + * @author 芋道源码
  23 + */
  24 +public class LlamaChatModelTests {
  25 +
  26 + private final OllamaApi ollamaApi = new OllamaApi(
  27 + "http://127.0.0.1:11434");
  28 + private final OllamaChatModel chatModel = new OllamaChatModel(ollamaApi,
  29 + OllamaOptions.create().withModel(OllamaModel.LLAMA3.getModelName()));
  30 +
  31 + @Test
  32 + @Disabled
  33 + public void testCall() {
  34 + // 准备参数
  35 + List<Message> messages = new ArrayList<>();
  36 + messages.add(new SystemMessage("你是一个优质的文言文作者，用文言文描述着各城市的人文风景。"));
  37 + messages.add(new UserMessage("1 + 1 = ？"));
  38 +
  39 + // 调用
  40 + ChatResponse response = chatModel.call(new Prompt(messages));
  41 + // 打印结果
  42 + System.out.println(response);
  43 + System.out.println(response.getResult().getOutput());
  44 + }
  45 +
  46 + @Test
  47 + @Disabled
  48 + public void testStream() {
  49 + // 准备参数
  50 + List<Message> messages = new ArrayList<>();
  51 + messages.add(new SystemMessage("你是一个优质的文言文作者，用文言文描述着各城市的人文风景。"));
  52 + messages.add(new UserMessage("1 + 1 = ？"));
  53 +
  54 + // 调用
  55 + Flux<ChatResponse> flux = chatModel.stream(new Prompt(messages));
  56 + // 打印结果
  57 + flux.doOnNext(response -> {
  58 +// System.out.println(response);
  59 + System.out.println(response.getResult().getOutput());
  60 + }).then().block();
  61 + }
  62 +
  63 +}
... ...