diff --git a/ai-rag-app/pom.xml b/ai-rag-app/pom.xml index 62aa8b5..7b22c7a 100644 --- a/ai-rag-app/pom.xml +++ b/ai-rag-app/pom.xml @@ -29,16 +29,16 @@ - - - - + + org.springframework.ai + spring-ai-tika-document-reader + - - - - - + + + org.springframework.ai + spring-ai-pgvector-store-spring-boot-starter + diff --git a/ai-rag-app/src/main/java/com/storm/dev/config/OllamaConfig.java b/ai-rag-app/src/main/java/com/storm/dev/config/OllamaConfig.java index 53299a2..db581be 100644 --- a/ai-rag-app/src/main/java/com/storm/dev/config/OllamaConfig.java +++ b/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); + } + + } \ No newline at end of file diff --git a/ai-rag-app/src/main/resources/application-dev.yml b/ai-rag-app/src/main/resources/application-dev.yml index c94cc65..f674300 100644 --- a/ai-rag-app/src/main/resources/application-dev.yml +++ b/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: diff --git a/ai-rag-app/src/test/java/com/storm/dev/text/RAGApiTest.java b/ai-rag-app/src/test/java/com/storm/dev/text/RAGApiTest.java new file mode 100644 index 0000000..50a51ba --- /dev/null +++ b/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 documents = reader.get(); + List 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 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 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)); + } +} diff --git a/ai-rag-app/src/test/resources/data/file.text b/ai-rag-app/src/test/resources/data/file.text new file mode 100644 index 0000000..8c39a93 --- /dev/null +++ b/ai-rag-app/src/test/resources/data/file.text @@ -0,0 +1 @@ +李永德，1999年12月31日出生，福建泉州人。 \ No newline at end of file