feat: rag的api
This commit is contained in:
@@ -27,14 +27,14 @@
|
||||
<!-- <groupId>org.springframework.ai</groupId>-->
|
||||
<!-- <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-pgvector-store</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-pgvector-store</artifactId>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.springframework.ai</groupId>
|
||||
<artifactId>spring-ai-ollama</artifactId>
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
package com.storm.dev.trigger.http;
|
||||
|
||||
import com.alibaba.fastjson.JSON;
|
||||
import com.storm.dev.api.IRAGService;
|
||||
import com.storm.dev.api.response.Response;
|
||||
import jakarta.annotation.Resource;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
import org.redisson.api.RList;
|
||||
import org.redisson.api.RedissonClient;
|
||||
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.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.web.bind.annotation.*;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
import reactor.core.publisher.Flux;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
/**
|
||||
* @author: lyd
|
||||
* @date: 2026/1/14 23:43
|
||||
*/
|
||||
@Slf4j
|
||||
@RestController()
|
||||
@CrossOrigin("*")
|
||||
@RequestMapping("/api/v1/rag/")
|
||||
public class RAGController implements IRAGService {
|
||||
@Resource
|
||||
private RedissonClient redissonClient;
|
||||
@Resource
|
||||
private OllamaChatClient ollamaChatClient;
|
||||
@Resource
|
||||
private TokenTextSplitter tokenTextSplitter;
|
||||
@Resource
|
||||
private SimpleVectorStore simpleVectorStore;
|
||||
@Resource
|
||||
private PgVectorStore pgVectorStore;
|
||||
@Override
|
||||
@RequestMapping(value = "query_rag_tag_list", method = RequestMethod.GET)
|
||||
public Response<List<String>> queryRagTagList() {
|
||||
RList<String> ragTag = redissonClient.getList("ragTag");
|
||||
return Response.<List<String>>builder()
|
||||
.code("0000")
|
||||
.info("调用成功")
|
||||
.data(ragTag)
|
||||
.build();
|
||||
}
|
||||
|
||||
@Override
|
||||
@RequestMapping(value = "file/upload", method = RequestMethod.POST, headers = "content-type=multipart/form-data")
|
||||
public Response<String> uploadFile(@RequestParam String ragTag, @RequestParam("file") List<MultipartFile> files) {
|
||||
log.info("上传知识库开始 {}", ragTag);
|
||||
for (MultipartFile file : files) {
|
||||
// 上传
|
||||
TikaDocumentReader reader = new TikaDocumentReader(file.getResource());
|
||||
List<Document> documents = reader.get();
|
||||
List<Document> documentSplitterList = tokenTextSplitter.apply(documents);
|
||||
// 打标
|
||||
documents.forEach(document -> document.getMetadata().put("knowledge", ragTag));
|
||||
documentSplitterList.forEach(document -> document.getMetadata().put("knowledge", ragTag));
|
||||
|
||||
pgVectorStore.accept(documentSplitterList);
|
||||
// 可以用MySQL存储
|
||||
RList<String> elements = redissonClient.getList("ragTag");
|
||||
if (!elements.contains(ragTag)){
|
||||
elements.add(ragTag);
|
||||
}
|
||||
log.info("上传完成!");
|
||||
}
|
||||
return Response.<String>builder().code("0000").info("调用成功").build();
|
||||
}
|
||||
|
||||
@Override
|
||||
@RequestMapping(value = "generate_stream_rag", method = RequestMethod.GET)
|
||||
public ChatResponse generateStreamRag(@RequestParam String model, @RequestParam String ragTag, @RequestParam String message) {
|
||||
log.info("用户选择模型:{},知识库:{},提问问题:{}", model, ragTag, 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 == '" + ragTag + "'");
|
||||
// 相似性搜索
|
||||
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);
|
||||
|
||||
// 提问
|
||||
// Flux<ChatResponse> chatResponse = ollamaChatClient.stream(new Prompt(messages, OllamaOptions.create().withModel(model)));
|
||||
ChatResponse call = ollamaChatClient.call(new Prompt(messages, OllamaOptions.create().withModel(model)));
|
||||
log.info("测试结果:{}", call);
|
||||
return call;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user