feat: rag测试

ai-rag-app/pom.xml

@@ -29,16 +29,16 @@
 <!--            <artifactId>spring-ai-openai-spring-boot-starter</artifactId>-->
 <!--        </dependency>-->
 
-<!--        <dependency>-->
-<!--            <groupId>org.springframework.ai</groupId>-->
-<!--            <artifactId>spring-ai-tika-document-reader</artifactId>-->
-<!--        </dependency>-->
+        <dependency>
+            <groupId>org.springframework.ai</groupId>
+            <artifactId>spring-ai-tika-document-reader</artifactId>
+        </dependency>
 
-<!--        &lt;!&ndash; 处理知识库：向量库 &ndash;&gt;-->
-<!--        <dependency>-->
-<!--            <groupId>org.springframework.ai</groupId>-->
-<!--            <artifactId>spring-ai-pgvector-store-spring-boot-starter</artifactId>-->
-<!--        </dependency>-->
+        <!-- 处理知识库：向量库 -->
+        <dependency>
+            <groupId>org.springframework.ai</groupId>
+            <artifactId>spring-ai-pgvector-store-spring-boot-starter</artifactId>
+        </dependency>
 
         <!-- 使用ollama的api -->
         <dependency>

ai-rag-app/src/main/java/com/storm/dev/config/OllamaConfig.java

@@ -1,10 +1,16 @@
 package com.storm.dev.config;
 
 import org.springframework.ai.ollama.OllamaChatClient;
+import org.springframework.ai.ollama.OllamaEmbeddingClient;
 import org.springframework.ai.ollama.api.OllamaApi;
+import org.springframework.ai.ollama.api.OllamaOptions;
+import org.springframework.ai.transformer.splitter.TokenTextSplitter;
+import org.springframework.ai.vectorstore.PgVectorStore;
+import org.springframework.ai.vectorstore.SimpleVectorStore;
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.jdbc.core.JdbcTemplate;
 
 /**
  * 注入OllamaApi、OllamaChatClient对象
@@ -24,4 +30,24 @@ public class OllamaConfig {
         return new OllamaChatClient(ollamaApi);
     }
 
+    @Bean
+    public TokenTextSplitter tokenTextSplitter() {
+        return new TokenTextSplitter();
+    }
+
+    @Bean
+    public SimpleVectorStore simpleVectorStore(OllamaApi ollamaApi) {
+        OllamaEmbeddingClient embeddingClient = new OllamaEmbeddingClient(ollamaApi);
+        embeddingClient.withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));
+        return new SimpleVectorStore(embeddingClient);
+    }
+
+    @Bean
+    public PgVectorStore pgVectorStore(OllamaApi ollamaApi, JdbcTemplate jdbcTemplate) {
+        OllamaEmbeddingClient embeddingClient = new OllamaEmbeddingClient(ollamaApi);
+        embeddingClient.withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));
+        return new PgVectorStore(jdbcTemplate, embeddingClient);
+    }
+
+
 }

ai-rag-app/src/main/resources/application-dev.yml

@@ -2,10 +2,19 @@ server:
   port: 8090
 
 spring:
+  datasource:
+    driver-class-name: org.postgresql.Driver
+    username: postgres
+    password: postgres
+    url: jdbc:postgresql://192.168.109.134:15432/ai-rag-knowledge
+    type: com.zaxxer.hikari.HikariDataSource
   ai:
     ollama:
       base-url: http://192.168.109.134:11434
-
+      embedding:
+        options:
+          num-batch: 512
+        model: nomic-embed-text
 # Redis
 redis:
   sdk:

ai-rag-app/src/test/java/com/storm/dev/text/RAGApiTest.java

@@ -0,0 +1,90 @@
+package com.storm.dev.text;
+
+import com.alibaba.fastjson.JSON;
+import jakarta.annotation.Resource;
+import lombok.extern.slf4j.Slf4j;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.springframework.ai.chat.ChatResponse;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.chat.prompt.SystemPromptTemplate;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.ollama.OllamaChatClient;
+import org.springframework.ai.ollama.api.OllamaApi;
+import org.springframework.ai.ollama.api.OllamaOptions;
+import org.springframework.ai.reader.tika.TikaDocumentReader;
+import org.springframework.ai.transformer.splitter.TokenTextSplitter;
+import org.springframework.ai.vectorstore.PgVectorStore;
+import org.springframework.ai.vectorstore.SearchRequest;
+import org.springframework.ai.vectorstore.SimpleVectorStore;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.test.context.junit4.SpringRunner;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/**
+ * @author: lyd
+ * @date: 2026/1/14 21:47
+ */
+@Slf4j
+@RunWith(SpringRunner.class)
+@SpringBootTest
+public class RAGApiTest {
+    @Resource
+    private OllamaChatClient ollamaChatClient;
+    @Resource
+    private TokenTextSplitter tokenTextSplitter;
+    @Resource
+    private SimpleVectorStore simpleVectorStore;
+    @Resource
+    private PgVectorStore pgVectorStore;
+
+    @Test
+    public void upload() {
+        // 上传
+        TikaDocumentReader reader = new TikaDocumentReader("./data/file.text");
+        List<Document> documents = reader.get();
+        List<Document> documentSplitterList = tokenTextSplitter.apply(documents);
+        // 打标
+        documents.forEach(document -> document.getMetadata().put("knowledge", "德德"));
+        documentSplitterList.forEach(document -> document.getMetadata().put("knowledge", "德德"));
+
+        pgVectorStore.accept(documentSplitterList);
+        log.info("上传完成!");
+    }
+
+    @Test
+    public void chat() {
+        // 构建提问
+        String message = "李永德，哪年出生的";
+
+        // 构建推理模板
+        String SYSTEM_PROMPT = """
+                Use the information from the DOCUMENTS section to provide accurate answers but act as if you knew this information innately.
+                If unsure, simply state that you don't know.
+                Another thing you need to note is that your reply must be in Chinese!
+                DOCUMENTS:
+                    {documents}
+                """;
+        // 读取向量库信息
+        SearchRequest request = SearchRequest.query(message).withTopK(5).withFilterExpression("knowledge == '德德'");
+        // 相似性搜索
+        List<Document> documents = pgVectorStore.similaritySearch(request);
+        String documentsCollectors = documents.stream().map(Document::getContent).collect(Collectors.joining());
+
+        // 推理：RAG
+        Message ragMessage = new SystemPromptTemplate(SYSTEM_PROMPT).createMessage(Map.of("documents", documentsCollectors));
+        ArrayList<Message> messages = new ArrayList<>();
+        messages.add(new UserMessage(message));
+        messages.add(ragMessage);
+
+        // 提问
+        ChatResponse chatResponse = ollamaChatClient.call(new Prompt(messages, OllamaOptions.create().withModel("deepseek-r1:7b")));
+        log.info("测试结果:{}", JSON.toJSONString(chatResponse));
+    }
+}

ai-rag-app/src/test/resources/data/file.text

@@ -0,0 +1 @@
+李永德，1999年12月31日出生，福建泉州人。