Эх сурвалжийг харах

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

YunaiV 5 сар өмнө
parent
commit
875cfa10c6

+ 65 - 0
yudao-module-ai/yudao-spring-boot-starter-ai/src/test/java/cn/iocoder/yudao/framework/ai/chat/OllamaChatModelTests.java

@@ -0,0 +1,65 @@
+package cn.iocoder.yudao.framework.ai.chat;
+
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.SystemMessage;
+import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.model.ChatResponse;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.ollama.OllamaChatModel;
+import org.springframework.ai.ollama.api.OllamaApi;
+import org.springframework.ai.ollama.api.OllamaOptions;
+import reactor.core.publisher.Flux;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * {@link OllamaChatModel} 集成测试
+ *
+ * @author 芋道源码
+ */
+public class OllamaChatModelTests {
+
+    private final OllamaChatModel chatModel = OllamaChatModel.builder()
+            .ollamaApi(new OllamaApi("http://127.0.0.1:11434")) // Ollama 服务地址
+            .defaultOptions(OllamaOptions.builder()
+//                    .model("qwen") // 模型(https://ollama.com/library/qwen)
+                    .model("deepseek-r1") // 模型(https://ollama.com/library/deepseek-r1)
+                    .build())
+            .build();
+
+    @Test
+    @Disabled
+    public void testCall() {
+        // 准备参数
+        List<Message> messages = new ArrayList<>();
+        messages.add(new SystemMessage("你是一个优质的文言文作者,用文言文描述着各城市的人文风景。"));
+        messages.add(new UserMessage("1 + 1 = ?"));
+
+        // 调用
+        ChatResponse response = chatModel.call(new Prompt(messages));
+        // 打印结果
+        System.out.println(response);
+        System.out.println(response.getResult().getOutput());
+    }
+
+    @Test
+    @Disabled
+    public void testStream() {
+        // 准备参数
+        List<Message> messages = new ArrayList<>();
+        messages.add(new SystemMessage("你是一个优质的文言文作者,用文言文描述着各城市的人文风景。"));
+        messages.add(new UserMessage("1 + 1 = ?"));
+
+        // 调用
+        Flux<ChatResponse> flux = chatModel.stream(new Prompt(messages));
+        // 打印结果
+        flux.doOnNext(response -> {
+//            System.out.println(response);
+            System.out.println(response.getResult().getOutput());
+        }).then().block();
+    }
+
+}