From 966c269681e56869161f971fee2c2b91fa1606e3 Mon Sep 17 00:00:00 2001 From: kazusa <409053122@qq.com> Date: Wed, 7 Jan 2026 18:55:13 +0800 Subject: [PATCH] init --- README.md | 53 +++ pom.xml | 97 ++++++ .../ai/assistant/AssistantApplication.java | 18 + .../assistant/config/AssistantProperties.java | 246 +++++++++++++ .../controller/IngestController.java | 33 ++ .../controller/OpenAIController.java | 143 ++++++++ .../ai/assistant/entity/VectorDocument.java | 71 ++++ .../mapper/VectorDocumentMapper.java | 22 ++ .../ai/assistant/model/AssistantAnswer.java | 56 +++ .../model/ChatCompletionRequest.java | 46 +++ .../model/ChatCompletionResponse.java | 136 ++++++++ .../ai/assistant/model/ChatMessage.java | 25 ++ .../assistant/service/AssistantService.java | 325 ++++++++++++++++++ .../ai/assistant/service/DomainFetcher.java | 205 +++++++++++ .../ai/assistant/service/DomainRouter.java | 106 ++++++ .../ai/assistant/service/TimeRangeParser.java | 70 ++++ .../assistant/service/VectorStoreService.java | 142 ++++++++ .../typehandler/JsonbTypeHandler.java | 41 +++ .../typehandler/VectorFloatTypeHandler.java | 73 ++++ src/main/resources/application.yml | 28 ++ .../data/sample_inventory_customer_a.jsonl | 2 + .../data/sample_orders_customer_a.jsonl | 2 + .../data/sample_returns_customer_a.jsonl | 2 + src/main/resources/domains.yml | 73 ++++ .../resources/mapper/VectorDocumentMapper.xml | 16 + 25 files changed, 2031 insertions(+) create mode 100644 README.md create mode 100644 pom.xml create mode 100644 src/main/java/cn/kazusa/ai/assistant/AssistantApplication.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/config/AssistantProperties.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/controller/IngestController.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/controller/OpenAIController.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/entity/VectorDocument.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/mapper/VectorDocumentMapper.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/model/AssistantAnswer.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/model/ChatCompletionRequest.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/model/ChatCompletionResponse.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/model/ChatMessage.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/service/AssistantService.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/service/DomainFetcher.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/service/DomainRouter.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/service/TimeRangeParser.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/service/VectorStoreService.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/typehandler/JsonbTypeHandler.java create mode 100644 src/main/java/cn/kazusa/ai/assistant/typehandler/VectorFloatTypeHandler.java create mode 100644 src/main/resources/application.yml create mode 100644 src/main/resources/data/sample_inventory_customer_a.jsonl create mode 100644 src/main/resources/data/sample_orders_customer_a.jsonl create mode 100644 src/main/resources/data/sample_returns_customer_a.jsonl create mode 100644 src/main/resources/domains.yml create mode 100644 src/main/resources/mapper/VectorDocumentMapper.xml diff --git a/README.md b/README.md new file mode 100644 index 0000000..3054961 --- /dev/null +++ b/README.md @@ -0,0 +1,53 @@ +# AI Assistant Java (Spring AI + pgvector) + +This module is an independent Spring Boot 3 service that provides an OpenAI-compatible API. + +## Requirements + +- Java 17+ +- PostgreSQL with pgvector +- Ollama (local inference) + +## Start pgvector (Docker) + +``` +docker run --name pgvector -e POSTGRES_PASSWORD=postgres -e POSTGRES_DB=assistant \ + -p 5432:5432 -d pgvector/pgvector:pg16 +``` + +## Start Ollama models + +``` +ollama pull qwen2.5:7b +ollama pull bge-m3 +``` + +## Configure + +- `src/main/resources/application.yml` +- `src/main/resources/domains.yml` + +## Run + +``` +./mvnw spring-boot:run +``` + +## OpenAI-compatible endpoint + +``` +POST http://127.0.0.1:8010/v1/chat/completions +``` + +Sample payload: + +``` +{ + "model": "local", + "tenant_id": "customer_a", + "stream": true, + "messages": [ + {"role": "user", "content": "订单有多少?"} + ] +} +``` diff --git a/pom.xml b/pom.xml new file mode 100644 index 0000000..db96263 --- /dev/null +++ b/pom.xml @@ -0,0 +1,97 @@ + + + 4.0.0 + + + org.springframework.boot + spring-boot-starter-parent + 3.2.5 + + + + cn.kazusa.ai + assistant + 0.1.0-SNAPSHOT + ai-assistant-java + Independent AI assistant service + + + 17 + 1.0.0-M5 + + + + + + org.springframework.ai + spring-ai-bom + ${spring-ai.version} + pom + import + + + + + + + org.springframework.boot + spring-boot-starter-web + + + com.baomidou + mybatis-plus-spring-boot3-starter + 3.5.5 + + + org.springframework.ai + spring-ai-ollama-spring-boot-starter + + + io.github.openfeign + feign-core + 13.2 + + + io.github.openfeign + feign-okhttp + 13.2 + + + org.postgresql + postgresql + + + org.springframework.boot + spring-boot-configuration-processor + true + + + org.springframework.boot + spring-boot-starter-validation + + + org.springframework.boot + spring-boot-starter-test + test + + + + + + spring-milestones + Spring Milestones + https://repo.spring.io/milestone + + + + + + + org.springframework.boot + spring-boot-maven-plugin + + + + diff --git a/src/main/java/cn/kazusa/ai/assistant/AssistantApplication.java b/src/main/java/cn/kazusa/ai/assistant/AssistantApplication.java new file mode 100644 index 0000000..0453801 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/AssistantApplication.java @@ -0,0 +1,18 @@ +package cn.kazusa.ai.assistant; + +import org.mybatis.spring.annotation.MapperScan; +import org.springframework.boot.SpringApplication; +import org.springframework.boot.autoconfigure.SpringBootApplication; +import org.springframework.boot.context.properties.ConfigurationPropertiesScan; + +/** + * @author zouzhiwen + */ +@SpringBootApplication +@ConfigurationPropertiesScan +@MapperScan("cn.kazusa.ai.assistant.mapper") +public class AssistantApplication { + public static void main(String[] args) { + SpringApplication.run(AssistantApplication.class, args); + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/config/AssistantProperties.java b/src/main/java/cn/kazusa/ai/assistant/config/AssistantProperties.java new file mode 100644 index 0000000..ca6ce7c --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/config/AssistantProperties.java @@ -0,0 +1,246 @@ +package cn.kazusa.ai.assistant.config; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.springframework.boot.context.properties.ConfigurationProperties; + +/** + * @author zouzhiwen + */ +@ConfigurationProperties(prefix = "assistant") +public class AssistantProperties { + private String defaultTenantId = "customer_a"; + private final Routing routing = new Routing(); + private final Vector vector = new Vector(); + private List domains = new ArrayList<>(); + + public String getDefaultTenantId() { + return defaultTenantId; + } + + public void setDefaultTenantId(String defaultTenantId) { + this.defaultTenantId = defaultTenantId; + } + + public Routing getRouting() { + return routing; + } + + public Vector getVector() { + return vector; + } + + public List getDomains() { + return domains; + } + + public void setDomains(List domains) { + this.domains = domains; + } + + public static class Routing { + private boolean useLlm = true; + private String fallbackDomain = "inventory"; + + public boolean isUseLlm() { + return useLlm; + } + + public void setUseLlm(boolean useLlm) { + this.useLlm = useLlm; + } + + public String getFallbackDomain() { + return fallbackDomain; + } + + public void setFallbackDomain(String fallbackDomain) { + this.fallbackDomain = fallbackDomain; + } + } + + public static class Vector { + private String table = "assistant_vectors"; + private int dimension = 1024; + + public String getTable() { + return table; + } + + public void setTable(String table) { + this.table = table; + } + + public int getDimension() { + return dimension; + } + + public void setDimension(int dimension) { + this.dimension = dimension; + } + } + + public static class DomainConfig { + private String id; + private String name; + private String queryMode = "live"; + private List keywords = new ArrayList<>(); + private List fields = new ArrayList<>(); + private Map statusMappings = new HashMap<>(); + private QuerySchema querySchema = new QuerySchema(); + private List sources = new ArrayList<>(); + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public String getName() { + return name; + } + + public void setName(String name) { + this.name = name; + } + + public String getQueryMode() { + return queryMode; + } + + public void setQueryMode(String queryMode) { + this.queryMode = queryMode; + } + + public List getKeywords() { + return keywords; + } + + public void setKeywords(List keywords) { + this.keywords = keywords; + } + + public List getFields() { + return fields; + } + + public void setFields(List fields) { + this.fields = fields; + } + + public Map getStatusMappings() { + return statusMappings; + } + + public void setStatusMappings(Map statusMappings) { + this.statusMappings = statusMappings; + } + + public QuerySchema getQuerySchema() { + return querySchema; + } + + public void setQuerySchema(QuerySchema querySchema) { + this.querySchema = querySchema; + } + + public List getSources() { + return sources; + } + + public void setSources(List sources) { + this.sources = sources; + } + } + + public static class QuerySchema { + private List allowedParams = new ArrayList<>(); + private Map> enums = new HashMap<>(); + + public List getAllowedParams() { + return allowedParams; + } + + public void setAllowedParams(List allowedParams) { + this.allowedParams = allowedParams; + } + + public Map> getEnums() { + return enums; + } + + public void setEnums(Map> enums) { + this.enums = enums; + } + } + + public static class DomainSource { + private String type; + private String url; + private String method = "POST"; + private String path; + private String itemsPath = ""; + private Map headers = new HashMap<>(); + private Map body = new HashMap<>(); + + public String getType() { + return type; + } + + public void setType(String type) { + this.type = type; + } + + public String getUrl() { + return url; + } + + public void setUrl(String url) { + this.url = url; + } + + public String getMethod() { + return method; + } + + public void setMethod(String method) { + this.method = method; + } + + public String getPath() { + return path; + } + + public void setPath(String path) { + this.path = path; + } + + public String getItemsPath() { + return itemsPath; + } + + public void setItemsPath(String itemsPath) { + this.itemsPath = itemsPath; + } + + public Map getHeaders() { + return headers; + } + + public void setHeaders(Map headers) { + this.headers = headers; + } + + public Map getBody() { + return body; + } + + public void setBody(Map body) { + this.body = body; + } + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/controller/IngestController.java b/src/main/java/cn/kazusa/ai/assistant/controller/IngestController.java new file mode 100644 index 0000000..448f593 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/controller/IngestController.java @@ -0,0 +1,33 @@ +package cn.kazusa.ai.assistant.controller; + +import cn.kazusa.ai.assistant.service.AssistantService; +import java.util.Map; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RestController; + +/** + * @author zouzhiwen + */ +@RestController +@RequestMapping("/ingest") +public class IngestController { + private final AssistantService assistantService; + + public IngestController(AssistantService assistantService) { + this.assistantService = assistantService; + } + + @PostMapping + public ResponseEntity ingest(@RequestBody Map payload) { + String tenantId = payload.get("tenant_id"); + String domainId = payload.get("domain_id"); + if (domainId == null || domainId.isEmpty()) { + return ResponseEntity.badRequest().body(Map.of("error", "domain_id required")); + } + assistantService.ingest(tenantId, domainId); + return ResponseEntity.ok(Map.of("ingested", true, "domain_id", domainId)); + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/controller/OpenAIController.java b/src/main/java/cn/kazusa/ai/assistant/controller/OpenAIController.java new file mode 100644 index 0000000..c195a73 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/controller/OpenAIController.java @@ -0,0 +1,143 @@ +package cn.kazusa.ai.assistant.controller; + +import cn.kazusa.ai.assistant.model.AssistantAnswer; +import cn.kazusa.ai.assistant.model.ChatCompletionRequest; +import cn.kazusa.ai.assistant.model.ChatCompletionResponse; +import cn.kazusa.ai.assistant.model.ChatMessage; +import cn.kazusa.ai.assistant.service.AssistantService; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.time.Instant; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.CrossOrigin; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RestController; +import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody; +import jakarta.servlet.http.HttpServletResponse; + +/** + * @author zouzhiwen + */ +@RestController +@RequestMapping("/v1") +@CrossOrigin +public class OpenAIController { + private final AssistantService assistantService; + private final ObjectMapper objectMapper; + + public OpenAIController(AssistantService assistantService, ObjectMapper objectMapper) { + this.assistantService = assistantService; + this.objectMapper = objectMapper; + } + + @PostMapping(value = "/chat/completions") + public Object chat(@RequestBody ChatCompletionRequest request, HttpServletResponse servletResponse) { + String question = extractQuestion(request.getMessages()); + if (question == null || question.isEmpty()) { + return ResponseEntity.badRequest().body(Map.of("error", "user message required")); + } + if (Boolean.TRUE.equals(request.getStream())) { + servletResponse.setContentType(MediaType.TEXT_EVENT_STREAM_VALUE); + return streamResponse(request, question); + } + AssistantAnswer answer = assistantService.answer(request.getTenantId(), question); + ChatCompletionResponse response = buildResponse(request.getModel(), answer.getAnswer()); + Map meta = new HashMap<>(); + meta.put("domain_id", answer.getDomainId()); + meta.put("route_reason", answer.getRouteReason()); + meta.put("sources", answer.getSources()); + meta.put("context", answer.getContext()); + meta.put("params", answer.getParams()); + response.setMeta(meta); + return ResponseEntity.ok(response); + } + + private StreamingResponseBody streamResponse(ChatCompletionRequest request, String question) { + return outputStream -> { + try { + writeEvent(outputStream, buildChunk("assistant", null, request.getModel(), true)); + assistantService.streamAnswer(request.getTenantId(), question, token -> { + try { + writeEvent(outputStream, buildChunk(null, token, request.getModel(), true)); + } catch (Exception ex) { + // 忽略单个分片的写入异常，保持流式响应不中断。 + } + }); + writeEvent(outputStream, buildChunk(null, null, request.getModel(), false)); + writeEvent(outputStream, "[DONE]"); + } catch (Exception ex) { + throw new RuntimeException("Stream write failed", ex); + } + }; + } + + private void writeEvent(OutputStream outputStream, Object payload) throws Exception { + String data = payload instanceof String ? (String) payload : objectMapper.writeValueAsString(payload); + String line = "data: " + data + "\n\n"; + outputStream.write(line.getBytes(StandardCharsets.UTF_8)); + outputStream.flush(); + } + + private ChatCompletionResponse buildResponse(String model, String content) { + ChatCompletionResponse response = new ChatCompletionResponse(); + response.setId("chatcmpl-local"); + response.setObject("chat.completion"); + response.setCreated(Instant.now().getEpochSecond()); + response.setModel(model == null ? "local" : model); + ChatCompletionResponse.Choice choice = new ChatCompletionResponse.Choice(); + choice.setIndex(0); + ChatMessage message = new ChatMessage(); + message.setRole("assistant"); + message.setContent(content); + choice.setMessage(message); + choice.setFinishReason("stop"); + response.setChoices(List.of(choice)); + ChatCompletionResponse.Usage usage = new ChatCompletionResponse.Usage(); + usage.setPromptTokens(0); + usage.setCompletionTokens(0); + usage.setTotalTokens(0); + response.setUsage(usage); + return response; + } + + private Map buildChunk(String role, String content, String model, boolean ongoing) { + Map response = new HashMap<>(); + response.put("id", "chatcmpl-local"); + response.put("object", "chat.completion.chunk"); + response.put("created", Instant.now().getEpochSecond()); + response.put("model", model == null ? "local" : model); + Map delta = new HashMap<>(); + if (role != null) { + delta.put("role", role); + } + if (content != null) { + delta.put("content", content); + } + Map choice = new HashMap<>(); + choice.put("index", 0); + choice.put("delta", delta); + choice.put("finish_reason", ongoing ? null : "stop"); + response.put("choices", List.of(choice)); + return response; + } + + private String extractQuestion(List messages) { + if (messages == null || messages.isEmpty()) { + return null; + } + for (int i = messages.size() - 1; i >= 0; i--) { + ChatMessage message = messages.get(i); + if ("user".equalsIgnoreCase(message.getRole())) { + return message.getContent(); + } + } + return null; + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/entity/VectorDocument.java b/src/main/java/cn/kazusa/ai/assistant/entity/VectorDocument.java new file mode 100644 index 0000000..83f6689 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/entity/VectorDocument.java @@ -0,0 +1,71 @@ +package cn.kazusa.ai.assistant.entity; + +import com.baomidou.mybatisplus.annotation.TableField; +import com.baomidou.mybatisplus.annotation.TableId; +import com.baomidou.mybatisplus.annotation.TableName; +import cn.kazusa.ai.assistant.typehandler.JsonbTypeHandler; +import cn.kazusa.ai.assistant.typehandler.VectorFloatTypeHandler; + +/** + * @author zouzhiwen + */ +@TableName(value = "assistant_vectors", autoResultMap = true) +public class VectorDocument { + @TableId + private String id; + private String tenantId; + private String domainId; + private String content; + @TableField(typeHandler = JsonbTypeHandler.class) + private String metadata; + @TableField(typeHandler = VectorFloatTypeHandler.class) + private float[] embedding; + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public String getTenantId() { + return tenantId; + } + + public void setTenantId(String tenantId) { + this.tenantId = tenantId; + } + + public String getDomainId() { + return domainId; + } + + public void setDomainId(String domainId) { + this.domainId = domainId; + } + + public String getContent() { + return content; + } + + public void setContent(String content) { + this.content = content; + } + + public String getMetadata() { + return metadata; + } + + public void setMetadata(String metadata) { + this.metadata = metadata; + } + + public float[] getEmbedding() { + return embedding; + } + + public void setEmbedding(float[] embedding) { + this.embedding = embedding; + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/mapper/VectorDocumentMapper.java b/src/main/java/cn/kazusa/ai/assistant/mapper/VectorDocumentMapper.java new file mode 100644 index 0000000..574b81c --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/mapper/VectorDocumentMapper.java @@ -0,0 +1,22 @@ +package cn.kazusa.ai.assistant.mapper; + +import cn.kazusa.ai.assistant.entity.VectorDocument; +import com.baomidou.mybatisplus.core.mapper.BaseMapper; +import org.apache.ibatis.annotations.Mapper; +import org.apache.ibatis.annotations.Param; + +import java.util.List; + +/** + * @author zouzhiwen + */ +@Mapper +public interface VectorDocumentMapper extends BaseMapper { + + List similaritySearch( + @Param("tenantId") String tenantId, + @Param("domainId") String domainId, + @Param("embedding") String embedding, + @Param("topK") int topK + ); +} diff --git a/src/main/java/cn/kazusa/ai/assistant/model/AssistantAnswer.java b/src/main/java/cn/kazusa/ai/assistant/model/AssistantAnswer.java new file mode 100644 index 0000000..90a4058 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/model/AssistantAnswer.java @@ -0,0 +1,56 @@ +package cn.kazusa.ai.assistant.model; + +import java.util.List; +import java.util.Map; + +/** + * @author zouzhiwen + */ +public class AssistantAnswer { + private final String answer; + private final List> sources; + private final String context; + private final String domainId; + private final String routeReason; + private final Map params; + + public AssistantAnswer( + String answer, + List> sources, + String context, + String domainId, + String routeReason, + Map params + ) { + this.answer = answer; + this.sources = sources; + this.context = context; + this.domainId = domainId; + this.routeReason = routeReason; + this.params = params; + } + + public String getAnswer() { + return answer; + } + + public List> getSources() { + return sources; + } + + public String getContext() { + return context; + } + + public String getDomainId() { + return domainId; + } + + public String getRouteReason() { + return routeReason; + } + + public Map getParams() { + return params; + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/model/ChatCompletionRequest.java b/src/main/java/cn/kazusa/ai/assistant/model/ChatCompletionRequest.java new file mode 100644 index 0000000..31b0ddd --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/model/ChatCompletionRequest.java @@ -0,0 +1,46 @@ +package cn.kazusa.ai.assistant.model; + +import java.util.ArrayList; +import java.util.List; + +/** + * @author zouzhiwen + */ +public class ChatCompletionRequest { + private String model; + private Boolean stream; + private String tenantId; + private List messages = new ArrayList<>(); + + public String getModel() { + return model; + } + + public void setModel(String model) { + this.model = model; + } + + public Boolean getStream() { + return stream; + } + + public void setStream(Boolean stream) { + this.stream = stream; + } + + public String getTenantId() { + return tenantId; + } + + public void setTenantId(String tenantId) { + this.tenantId = tenantId; + } + + public List getMessages() { + return messages; + } + + public void setMessages(List messages) { + this.messages = messages; + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/model/ChatCompletionResponse.java b/src/main/java/cn/kazusa/ai/assistant/model/ChatCompletionResponse.java new file mode 100644 index 0000000..960c8dd --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/model/ChatCompletionResponse.java @@ -0,0 +1,136 @@ +package cn.kazusa.ai.assistant.model; + +import com.fasterxml.jackson.annotation.JsonInclude; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; + +/** + * @author zouzhiwen + */ +@JsonInclude(JsonInclude.Include.NON_NULL) +public class ChatCompletionResponse { + private String id; + private String object; + private long created; + private String model; + private List choices = new ArrayList<>(); + private Usage usage; + private Map meta; + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public String getObject() { + return object; + } + + public void setObject(String object) { + this.object = object; + } + + public long getCreated() { + return created; + } + + public void setCreated(long created) { + this.created = created; + } + + public String getModel() { + return model; + } + + public void setModel(String model) { + this.model = model; + } + + public List getChoices() { + return choices; + } + + public void setChoices(List choices) { + this.choices = choices; + } + + public Usage getUsage() { + return usage; + } + + public void setUsage(Usage usage) { + this.usage = usage; + } + + public Map getMeta() { + return meta; + } + + public void setMeta(Map meta) { + this.meta = meta; + } + + public static class Choice { + private int index; + private ChatMessage message; + private String finishReason; + + public int getIndex() { + return index; + } + + public void setIndex(int index) { + this.index = index; + } + + public ChatMessage getMessage() { + return message; + } + + public void setMessage(ChatMessage message) { + this.message = message; + } + + public String getFinishReason() { + return finishReason; + } + + public void setFinishReason(String finishReason) { + this.finishReason = finishReason; + } + } + + public static class Usage { + private int promptTokens; + private int completionTokens; + private int totalTokens; + + public int getPromptTokens() { + return promptTokens; + } + + public void setPromptTokens(int promptTokens) { + this.promptTokens = promptTokens; + } + + public int getCompletionTokens() { + return completionTokens; + } + + public void setCompletionTokens(int completionTokens) { + this.completionTokens = completionTokens; + } + + public int getTotalTokens() { + return totalTokens; + } + + public void setTotalTokens(int totalTokens) { + this.totalTokens = totalTokens; + } + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/model/ChatMessage.java b/src/main/java/cn/kazusa/ai/assistant/model/ChatMessage.java new file mode 100644 index 0000000..a2cbc88 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/model/ChatMessage.java @@ -0,0 +1,25 @@ +package cn.kazusa.ai.assistant.model; + +/** + * @author zouzhiwen + */ +public class ChatMessage { + private String role; + private String content; + + public String getRole() { + return role; + } + + public void setRole(String role) { + this.role = role; + } + + public String getContent() { + return content; + } + + public void setContent(String content) { + this.content = content; + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/service/AssistantService.java b/src/main/java/cn/kazusa/ai/assistant/service/AssistantService.java new file mode 100644 index 0000000..f9228ef --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/service/AssistantService.java @@ -0,0 +1,325 @@ +package cn.kazusa.ai.assistant.service; + +import cn.kazusa.ai.assistant.config.AssistantProperties; +import cn.kazusa.ai.assistant.model.AssistantAnswer; +import cn.kazusa.ai.assistant.service.DomainRouter.RouteResult; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.stereotype.Service; + +/** + * @author zouzhiwen + */ +@Service +public class AssistantService { + private final AssistantProperties properties; + private final DomainRouter domainRouter; + private final DomainFetcher domainFetcher; + private final VectorStoreService vectorStoreService; + private final ChatModel chatModel; + private final ObjectMapper objectMapper; + + public AssistantService( + AssistantProperties properties, + DomainRouter domainRouter, + DomainFetcher domainFetcher, + VectorStoreService vectorStoreService, + ChatModel chatModel, + ObjectMapper objectMapper + ) { + this.properties = properties; + this.domainRouter = domainRouter; + this.domainFetcher = domainFetcher; + this.vectorStoreService = vectorStoreService; + this.chatModel = chatModel; + this.objectMapper = objectMapper; + } + + public AssistantAnswer answer(String tenantId, String question) { + ContextResult contextResult = buildContext(tenantId, question); + String answer = callChatModel(contextResult.getContextText(), question); + return new AssistantAnswer( + answer, + contextResult.getSources(), + contextResult.getContextText(), + contextResult.getDomainId(), + contextResult.getRouteReason(), + contextResult.getParams() + ); + } + + public void streamAnswer(String tenantId, String question, java.util.function.Consumer onToken) { + ContextResult contextResult = buildContext(tenantId, question); + Prompt prompt = buildPrompt(contextResult.getContextText(), question); + chatModel.stream(prompt).toStream().forEach(response -> { + String token = response.getResult().getOutput().getContent(); + if (token != null && !token.isEmpty()) { + onToken.accept(token); + } + }); + } + + public void ingest(String tenantId, String domainId) { + AssistantProperties.DomainConfig domain = findDomain(domainId); + String resolvedTenant = tenantId == null || tenantId.isEmpty() ? properties.getDefaultTenantId() : tenantId; + List> records = domainFetcher.fetchRecords(domain, resolvedTenant, "", Map.of()); + List documents = toContextRecords(records, resolvedTenant, domain); + for (VectorStoreService.DocumentRecord doc : documents) { + float[] embedding = vectorStoreService.embed(doc.getContent()); + VectorStoreService.DocumentRecord enriched = new VectorStoreService.DocumentRecord( + doc.getId(), + doc.getContent(), + doc.getMetadata(), + embedding + ); + vectorStoreService.upsert(resolvedTenant, domainId, List.of(enriched)); + } + } + + private AssistantProperties.DomainConfig findDomain(String domainId) { + Optional match = properties.getDomains().stream() + .filter(domain -> domain.getId().equals(domainId)) + .findFirst(); + return match.orElseThrow(() -> new IllegalArgumentException("Unknown domain: " + domainId)); + } + + private Map llmExtractParams(String question, AssistantProperties.DomainConfig domain) { + if (domain.getQuerySchema().getAllowedParams().isEmpty()) { + return Map.of(); + } + // 将可用参数范围明确给模型，减少无关字段与幻觉。 + String prompt = "Extract query parameters as JSON using ONLY the allowed fields. " + + "If a field is not mentioned, omit it. Return JSON only.\n\n" + + "Allowed fields: " + domain.getQuerySchema().getAllowedParams() + "\n" + + "Enums: " + domain.getQuerySchema().getEnums() + "\n\n" + + "Question:\n" + question; + Prompt chatPrompt = new Prompt(List.of(new UserMessage(prompt))); + String content = chatModel.call(chatPrompt).getResult().getOutput().getContent(); + try { + Map raw = objectMapper.readValue(content, Map.class); + Map filtered = new HashMap<>(); + for (String key : domain.getQuerySchema().getAllowedParams()) { + if (raw.containsKey(key)) { + filtered.put(key, raw.get(key)); + } + } + return filtered; + } catch (Exception ex) { + return Map.of(); + } + } + + private Map normalizeParams( + Map params, + AssistantProperties.DomainConfig domain, + String question + ) { + // 对抽取参数进行校验并补充时间范围，避免非法值进入业务 API。 + Map normalized = new HashMap<>(); + Map> enums = domain.getQuerySchema().getEnums(); + for (Map.Entry entry : params.entrySet()) { + String key = entry.getKey(); + Object value = entry.getValue(); + if (value == null) { + continue; + } + if (enums.containsKey(key)) { + List allowed = enums.get(key); + if (value instanceof List) { + List filtered = new ArrayList<>(); + for (Object item : (List) value) { + if (allowed.contains(String.valueOf(item))) { + filtered.add(String.valueOf(item)); + } + } + if (!filtered.isEmpty()) { + normalized.put(key, String.join(",", filtered)); + } + } else if (allowed.contains(String.valueOf(value))) { + normalized.put(key, String.valueOf(value)); + } + continue; + } + if ("time_range".equals(key) && value instanceof Map) { + Map range = (Map) value; + Object from = range.get("from"); + Object to = range.get("to"); + if (from != null && to != null) { + normalized.put("time_from", String.valueOf(from)); + normalized.put("time_to", String.valueOf(to)); + } + continue; + } + normalized.put(key, String.valueOf(value)); + } + if (!normalized.containsKey("time_from") || !normalized.containsKey("time_to")) { + Map resolved = TimeRangeParser.resolve(question); + if (resolved.containsKey("time_from") && resolved.containsKey("time_to")) { + normalized.putIfAbsent("time_from", resolved.get("time_from")); + normalized.putIfAbsent("time_to", resolved.get("time_to")); + } + } + return normalized; + } + + private List toContextRecords( + List> records, + String tenantId, + AssistantProperties.DomainConfig domain + ) { + List result = new ArrayList<>(); + for (Map record : records) { + String content = renderRecord(record, domain); + Map metadata = new HashMap<>(); + metadata.put("tenant_id", tenantId); + metadata.put("domain_id", domain.getId()); + metadata.put("source", domain.getId()); + metadata.put("record_id", recordId(record)); + result.add(new VectorStoreService.DocumentRecord(recordId(record), content, metadata, new float[0])); + } + return result; + } + + private String renderRecord(Map record, AssistantProperties.DomainConfig domain) { + List fields = domain.getFields(); + if (fields.isEmpty()) { + return record.toString(); + } + StringBuilder builder = new StringBuilder(); + for (String field : fields) { + Object value = record.get(field); + if (value == null) { + continue; + } + String displayValue = String.valueOf(value); + if ("status".equals(field) && domain.getStatusMappings().containsKey(displayValue)) { + displayValue = displayValue + " (" + domain.getStatusMappings().get(displayValue) + ")"; + } + builder.append(field).append(": ").append(displayValue).append("\n"); + } + return builder.toString().trim(); + } + + private String recordId(Map record) { + Object id = record.get("id"); + if (id == null) { + id = record.get("record_id"); + } + if (id == null) { + id = record.get("sku"); + } + if (id == null) { + id = record.get("order_no"); + } + if (id == null) { + id = record.get("return_no"); + } + return id == null ? "unknown" : String.valueOf(id); + } + + private String buildContextText(List records) { + StringBuilder builder = new StringBuilder(); + for (VectorStoreService.DocumentRecord record : records) { + Map meta = record.getMetadata(); + builder.append("source: ").append(meta.get("source")) + .append(" record_id: ").append(meta.get("record_id")) + .append("\n").append(record.getContent()).append("\n\n"); + } + return builder.toString().trim(); + } + + private Prompt buildPrompt(String context, String question) { + String systemPrompt = "You are a warehouse assistant. Answer using ONLY the provided context. " + + "If the answer is not in the context, say you do not have enough data. " + + "Keep answers concise and factual."; + String userPrompt = "Context:\n" + context + "\n\nQuestion:\n" + question; + return new Prompt(List.of(new UserMessage(systemPrompt), new UserMessage(userPrompt))); + } + + private String callChatModel(String context, String question) { + Prompt prompt = buildPrompt(context, question); + return chatModel.call(prompt).getResult().getOutput().getContent(); + } + + private ContextResult buildContext(String tenantId, String question) { + String resolvedTenant = tenantId == null || tenantId.isEmpty() ? properties.getDefaultTenantId() : tenantId; + RouteResult routeResult = domainRouter.route(question); + AssistantProperties.DomainConfig domain = findDomain(routeResult.getDomainId()); + Map extractedParams = new HashMap<>(); + List contextRecords = new ArrayList<>(); + + if ("live".equalsIgnoreCase(domain.getQueryMode())) { + extractedParams = llmExtractParams(question, domain); + Map normalized = normalizeParams(extractedParams, domain, question); + List> records = domainFetcher.fetchRecords(domain, resolvedTenant, question, normalized); + contextRecords = toContextRecords(records, resolvedTenant, domain); + if (contextRecords.isEmpty()) { + contextRecords = vectorStoreService.similaritySearch(resolvedTenant, domain.getId(), question, 4); + } + } else { + contextRecords = vectorStoreService.similaritySearch(resolvedTenant, domain.getId(), question, 4); + } + + String contextText = buildContextText(contextRecords); + List> sources = new ArrayList<>(); + for (VectorStoreService.DocumentRecord record : contextRecords) { + sources.add(record.getMetadata()); + } + return new ContextResult( + domain.getId(), + routeResult.getReason(), + contextText, + sources, + extractedParams + ); + } + + private static class ContextResult { + private final String domainId; + private final String routeReason; + private final String contextText; + private final List> sources; + private final Map params; + + private ContextResult( + String domainId, + String routeReason, + String contextText, + List> sources, + Map params + ) { + this.domainId = domainId; + this.routeReason = routeReason; + this.contextText = contextText; + this.sources = sources; + this.params = params; + } + + public String getDomainId() { + return domainId; + } + + public String getRouteReason() { + return routeReason; + } + + public String getContextText() { + return contextText; + } + + public List> getSources() { + return sources; + } + + public Map getParams() { + return params; + } + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/service/DomainFetcher.java b/src/main/java/cn/kazusa/ai/assistant/service/DomainFetcher.java new file mode 100644 index 0000000..bc9eded --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/service/DomainFetcher.java @@ -0,0 +1,205 @@ +package cn.kazusa.ai.assistant.service; + +import cn.kazusa.ai.assistant.config.AssistantProperties; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.io.BufferedReader; +import java.io.InputStreamReader; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Locale; +import java.util.concurrent.TimeUnit; +import feign.Request; +import feign.Response; +import feign.okhttp.OkHttpClient; +import org.springframework.core.io.Resource; +import org.springframework.core.io.ResourceLoader; +import org.springframework.stereotype.Component; + +/** + * @author zouzhiwen + */ +@Component +public class DomainFetcher { + private final OkHttpClient client; + private final ObjectMapper objectMapper; + private final ResourceLoader resourceLoader; + + public DomainFetcher(ObjectMapper objectMapper, ResourceLoader resourceLoader) { + this.client = new OkHttpClient(); + this.objectMapper = objectMapper; + this.resourceLoader = resourceLoader; + } + + public List> fetchRecords( + AssistantProperties.DomainConfig domain, + String tenantId, + String question, + Map params + ) { + // 每个业务域允许配置多个数据源，按顺序汇总。 + List> result = new ArrayList<>(); + for (AssistantProperties.DomainSource source : domain.getSources()) { + if ("jsonl".equalsIgnoreCase(source.getType())) { + result.addAll(loadJsonl(source.getPath())); + } else if ("api".equalsIgnoreCase(source.getType())) { + result.addAll(loadApi(source, tenantId, question, params)); + } + } + return result; + } + + private List> loadJsonl(String path) { + if (path == null || path.isEmpty()) { + return List.of(); + } + Resource resource = resourceLoader.getResource(path); + if (!resource.exists()) { + resource = resourceLoader.getResource("classpath:" + path); + } + List> items = new ArrayList<>(); + try (BufferedReader reader = new BufferedReader( + new InputStreamReader(resource.getInputStream(), StandardCharsets.UTF_8))) { + String line; + while ((line = reader.readLine()) != null) { + String trimmed = line.trim(); + if (trimmed.isEmpty()) { + continue; + } + items.add(objectMapper.readValue(trimmed, Map.class)); + } + } catch (Exception ex) { + return List.of(); + } + return items; + } + + private List> loadApi( + AssistantProperties.DomainSource source, + String tenantId, + String question, + Map params + ) { + Map replacements = new HashMap<>(); + replacements.put("tenant_id", tenantId); + replacements.put("question", question); + if (params != null) { + replacements.putAll(params); + } + String url = replacePlaceholders(source.getUrl(), replacements); + Map body = replacePlaceholders(source.getBody(), replacements); + body = stripUnresolved(body); + Request.HttpMethod method = Request.HttpMethod.valueOf(source.getMethod().toUpperCase(Locale.ROOT)); + byte[] bodyBytes = method == Request.HttpMethod.GET ? null : serializeBody(body); + Map> headers = new HashMap<>(); + for (Map.Entry header : source.getHeaders().entrySet()) { + headers.put(header.getKey(), List.of(header.getValue())); + } + Request request = Request.create(method, url, headers, bodyBytes, StandardCharsets.UTF_8, null); + try (Response response = client.execute(request, new Request.Options(30, TimeUnit.SECONDS, 30, TimeUnit.SECONDS, true))) { + if (response.status() >= 400 || response.body() == null) { + return List.of(); + } + JsonNode json = objectMapper.readTree(response.body().asInputStream()); + return extractItems(json, source.getItemsPath()); + } catch (Exception ex) { + return List.of(); + } + } + + private String replacePlaceholders(String template, Map replacements) { + if (template == null) { + return ""; + } + String result = template; + for (Map.Entry entry : replacements.entrySet()) { + result = result.replace("{" + entry.getKey() + "}", entry.getValue()); + } + return result; + } + + private Map replacePlaceholders(Map template, Map replacements) { + Map result = new HashMap<>(); + for (Map.Entry entry : template.entrySet()) { + result.put(entry.getKey(), replaceValue(entry.getValue(), replacements)); + } + return result; + } + + private Object replaceValue(Object value, Map replacements) { + if (value instanceof String) { + return replacePlaceholders((String) value, replacements); + } + if (value instanceof Map) { + Map nested = new HashMap<>(); + Map input = (Map) value; + for (Map.Entry entry : input.entrySet()) { + nested.put(String.valueOf(entry.getKey()), replaceValue(entry.getValue(), replacements)); + } + return nested; + } + if (value instanceof List) { + List nested = new ArrayList<>(); + for (Object item : (List) value) { + nested.add(replaceValue(item, replacements)); + } + return nested; + } + return value; + } + + private Map stripUnresolved(Map body) { + Map cleaned = new HashMap<>(); + for (Map.Entry entry : body.entrySet()) { + Object value = entry.getValue(); + if (value instanceof String && ((String) value).contains("{") && ((String) value).contains("}")) { + continue; + } + if (value instanceof Map) { + Map nested = stripUnresolved((Map) value); + if (!nested.isEmpty()) { + cleaned.put(entry.getKey(), nested); + } + continue; + } + cleaned.put(entry.getKey(), value); + } + return cleaned; + } + + private byte[] serializeBody(Map body) { + try { + return objectMapper.writeValueAsBytes(body); + } catch (Exception ex) { + return "{}".getBytes(StandardCharsets.UTF_8); + } + } + + private List> extractItems(JsonNode data, String itemsPath) { + if (data == null || data.isNull()) { + return List.of(); + } + JsonNode current = data; + if (itemsPath != null && !itemsPath.isEmpty()) { + String[] parts = itemsPath.split("\\."); + for (String part : parts) { + if (part.isEmpty()) { + continue; + } + current = current.path(part); + } + } + if (current == null || !current.isArray()) { + return List.of(); + } + List> items = new ArrayList<>(); + for (JsonNode node : current) { + items.add(objectMapper.convertValue(node, Map.class)); + } + return items; + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/service/DomainRouter.java b/src/main/java/cn/kazusa/ai/assistant/service/DomainRouter.java new file mode 100644 index 0000000..33f8537 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/service/DomainRouter.java @@ -0,0 +1,106 @@ +package cn.kazusa.ai.assistant.service; + +import cn.kazusa.ai.assistant.config.AssistantProperties; +import cn.kazusa.ai.assistant.model.AssistantAnswer; +import java.util.List; +import java.util.Locale; +import java.util.Optional; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.stereotype.Component; + +/** + * @author zouzhiwen + */ +@Component +public class DomainRouter { + private final AssistantProperties properties; + private final ChatModel chatModel; + + public DomainRouter(AssistantProperties properties, ChatModel chatModel) { + this.properties = properties; + this.chatModel = chatModel; + } + + public RouteResult route(String question) { + Optional ruleResult = matchByKeyword(question); + if (ruleResult.isPresent()) { + return ruleResult.get(); + } + if (!properties.getRouting().isUseLlm()) { + return new RouteResult(properties.getRouting().getFallbackDomain(), "fallback"); + } + String llmDomain = classifyByLlm(question); + if (llmDomain != null && !llmDomain.isEmpty()) { + return new RouteResult(llmDomain, "llm_classification"); + } + return new RouteResult(properties.getRouting().getFallbackDomain(), "fallback"); + } + + private Optional matchByKeyword(String question) { + String normalized = question == null ? "" : question.toLowerCase(Locale.ROOT); + int bestScore = 0; + String bestDomain = null; + for (AssistantProperties.DomainConfig domain : properties.getDomains()) { + int score = 0; + for (String keyword : domain.getKeywords()) { + if (keyword != null && !keyword.isEmpty() && normalized.contains(keyword.toLowerCase(Locale.ROOT))) { + score++; + } + } + if (score > bestScore) { + bestScore = score; + bestDomain = domain.getId(); + } + } + if (bestDomain == null) { + return Optional.empty(); + } + return Optional.of(new RouteResult(bestDomain, "keyword_match:" + bestScore)); + } + + private String classifyByLlm(String question) { + StringBuilder domainList = new StringBuilder(); + for (AssistantProperties.DomainConfig domain : properties.getDomains()) { + domainList.append("- ") + .append(domain.getId()) + .append(": ") + .append(domain.getName()) + .append("\n"); + } + String prompt = "Select the best domain for the question. " + + "Return only the domain id from the list.\n\n" + + "Domains:\n" + domainList + + "\nQuestion:\n" + question; + Prompt chatPrompt = new Prompt(List.of(new UserMessage(prompt))); + String content = chatModel.call(chatPrompt).getResult().getOutput().getContent(); + if (content == null) { + return null; + } + for (AssistantProperties.DomainConfig domain : properties.getDomains()) { + if (content.contains(domain.getId())) { + return domain.getId(); + } + } + return null; + } + + public static class RouteResult { + private final String domainId; + private final String reason; + + public RouteResult(String domainId, String reason) { + this.domainId = domainId; + this.reason = reason; + } + + public String getDomainId() { + return domainId; + } + + public String getReason() { + return reason; + } + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/service/TimeRangeParser.java b/src/main/java/cn/kazusa/ai/assistant/service/TimeRangeParser.java new file mode 100644 index 0000000..0f255ab --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/service/TimeRangeParser.java @@ -0,0 +1,70 @@ +package cn.kazusa.ai.assistant.service; + +import java.time.DayOfWeek; +import java.time.LocalDate; +import java.time.temporal.TemporalAdjusters; +import java.util.HashMap; +import java.util.Locale; +import java.util.Map; + +/** + * @author zouzhiwen + */ +public final class TimeRangeParser { + private TimeRangeParser() { + } + + public static Map resolve(String question) { + // 只处理常见中文相对时间表达，无法识别时返回空。 + Map result = new HashMap<>(); + if (question == null || question.isEmpty()) { + return result; + } + String q = question.replace(" ", ""); + LocalDate today = LocalDate.now(); + if (q.contains("近一周") || q.contains("最近一周") || q.contains("过去一周")) { + return range(today.minusDays(7), today); + } + if (q.contains("近一个月") || q.contains("最近一个月") || q.contains("过去一个月")) { + return range(today.minusDays(30), today); + } + if (q.contains("近30天") || q.contains("最近30天") || q.contains("过去30天")) { + return range(today.minusDays(30), today); + } + if (q.contains("本周")) { + LocalDate start = today.with(TemporalAdjusters.previousOrSame(DayOfWeek.MONDAY)); + return range(start, today); + } + if (q.contains("上周")) { + LocalDate start = today.with(TemporalAdjusters.previousOrSame(DayOfWeek.MONDAY)).minusWeeks(1); + LocalDate end = start.plusDays(6); + return range(start, end); + } + if (q.contains("本月")) { + LocalDate start = today.withDayOfMonth(1); + return range(start, today); + } + if (q.contains("上月")) { + LocalDate start = today.minusMonths(1).withDayOfMonth(1); + LocalDate end = start.plusMonths(1).minusDays(1); + return range(start, end); + } + if (q.contains("今年")) { + LocalDate start = LocalDate.of(today.getYear(), 1, 1); + return range(start, today); + } + if (q.contains("去年")) { + LocalDate start = LocalDate.of(today.getYear() - 1, 1, 1); + LocalDate end = LocalDate.of(today.getYear() - 1, 12, 31); + return range(start, end); + } + return result; + } + + private static Map range(LocalDate from, LocalDate to) { + Map result = new HashMap<>(); + result.put("time_from", from.toString()); + result.put("time_to", to.toString()); + return result; + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/service/VectorStoreService.java b/src/main/java/cn/kazusa/ai/assistant/service/VectorStoreService.java new file mode 100644 index 0000000..39479a0 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/service/VectorStoreService.java @@ -0,0 +1,142 @@ +package cn.kazusa.ai.assistant.service; + +import cn.kazusa.ai.assistant.config.AssistantProperties; +import cn.kazusa.ai.assistant.entity.VectorDocument; +import cn.kazusa.ai.assistant.mapper.VectorDocumentMapper; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.stereotype.Component; + +/** + * @author zouzhiwen + */ +@Component +public class VectorStoreService { + private final VectorDocumentMapper vectorDocumentMapper; + private final JdbcTemplate jdbcTemplate; + private final EmbeddingModel embeddingModel; + private final ObjectMapper objectMapper; + private final AssistantProperties properties; + + public VectorStoreService( + VectorDocumentMapper vectorDocumentMapper, + JdbcTemplate jdbcTemplate, + EmbeddingModel embeddingModel, + ObjectMapper objectMapper, + AssistantProperties properties + ) { + this.vectorDocumentMapper = vectorDocumentMapper; + this.jdbcTemplate = jdbcTemplate; + this.embeddingModel = embeddingModel; + this.objectMapper = objectMapper; + this.properties = properties; + initSchema(); + } + + public void upsert(String tenantId, String domainId, List records) { + for (DocumentRecord record : records) { + VectorDocument doc = new VectorDocument(); + doc.setId(record.getId()); + doc.setTenantId(tenantId); + doc.setDomainId(domainId); + doc.setContent(record.getContent()); + try { + doc.setMetadata(objectMapper.writeValueAsString(record.getMetadata())); + } catch (Exception e) { + doc.setMetadata("{}"); + } + doc.setEmbedding(record.getEmbedding()); + + VectorDocument existing = vectorDocumentMapper.selectById(record.getId()); + if (existing != null) { + vectorDocumentMapper.updateById(doc); + } else { + vectorDocumentMapper.insert(doc); + } + } + } + + public List similaritySearch(String tenantId, String domainId, String query, int topK) { + float[] embedding = embeddingModel.embed(query); + String embeddingStr = arrayToString(embedding); + List docs = vectorDocumentMapper.similaritySearch(tenantId, domainId, embeddingStr, topK); + List result = new ArrayList<>(); + for (VectorDocument doc : docs) { + Map metadata; + try { + metadata = objectMapper.readValue(doc.getMetadata(), Map.class); + } catch (Exception e) { + metadata = new HashMap<>(); + } + result.add(new DocumentRecord(doc.getId(), doc.getContent(), metadata, new float[0])); + } + return result; + } + + public float[] embed(String text) { + return embeddingModel.embed(text); + } + + private String arrayToString(float[] values) { + StringBuilder builder = new StringBuilder(); + builder.append('['); + for (int i = 0; i < values.length; i++) { + if (i > 0) { + builder.append(','); + } + builder.append(values[i]); + } + builder.append(']'); + return builder.toString(); + } + + private void initSchema() { + String table = properties.getVector().getTable(); + int dimension = properties.getVector().getDimension(); + jdbcTemplate.execute("CREATE EXTENSION IF NOT EXISTS vector"); + jdbcTemplate.execute("CREATE TABLE IF NOT EXISTS " + table + " (" + + "id TEXT PRIMARY KEY," + + "tenant_id TEXT NOT NULL," + + "domain_id TEXT NOT NULL," + + "content TEXT NOT NULL," + + "metadata JSONB NOT NULL," + + "embedding vector(" + dimension + ")" + + ")"); + jdbcTemplate.execute("CREATE INDEX IF NOT EXISTS idx_" + table + "_tenant_domain ON " + table + " (tenant_id, domain_id)"); + } + + public static class DocumentRecord { + private final String id; + private final String content; + private final Map metadata; + private final float[] embedding; + + public DocumentRecord(String id, String content, Map metadata, float[] embedding) { + this.id = id; + this.content = content; + this.metadata = metadata; + this.embedding = embedding; + } + + public String getId() { + return id; + } + + public String getContent() { + return content; + } + + public Map getMetadata() { + return metadata; + } + + public float[] getEmbedding() { + return embedding; + } + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/typehandler/JsonbTypeHandler.java b/src/main/java/cn/kazusa/ai/assistant/typehandler/JsonbTypeHandler.java new file mode 100644 index 0000000..419cf42 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/typehandler/JsonbTypeHandler.java @@ -0,0 +1,41 @@ +package cn.kazusa.ai.assistant.typehandler; + +import org.apache.ibatis.type.BaseTypeHandler; +import org.apache.ibatis.type.JdbcType; +import org.apache.ibatis.type.MappedTypes; +import org.postgresql.util.PGobject; + +import java.sql.CallableStatement; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; + +/** + * @author zouzhiwen + */ +@MappedTypes(String.class) +public class JsonbTypeHandler extends BaseTypeHandler { + + @Override + public void setNonNullParameter(PreparedStatement ps, int i, String parameter, JdbcType jdbcType) throws SQLException { + PGobject pgObject = new PGobject(); + pgObject.setType("jsonb"); + pgObject.setValue(parameter); + ps.setObject(i, pgObject); + } + + @Override + public String getNullableResult(ResultSet rs, String columnName) throws SQLException { + return rs.getString(columnName); + } + + @Override + public String getNullableResult(ResultSet rs, int columnIndex) throws SQLException { + return rs.getString(columnIndex); + } + + @Override + public String getNullableResult(CallableStatement cs, int columnIndex) throws SQLException { + return cs.getString(columnIndex); + } +} diff --git a/src/main/java/cn/kazusa/ai/assistant/typehandler/VectorFloatTypeHandler.java b/src/main/java/cn/kazusa/ai/assistant/typehandler/VectorFloatTypeHandler.java new file mode 100644 index 0000000..cbe61e0 --- /dev/null +++ b/src/main/java/cn/kazusa/ai/assistant/typehandler/VectorFloatTypeHandler.java @@ -0,0 +1,73 @@ +package cn.kazusa.ai.assistant.typehandler; + +import org.apache.ibatis.type.BaseTypeHandler; +import org.apache.ibatis.type.JdbcType; +import org.apache.ibatis.type.MappedTypes; +import org.postgresql.util.PGobject; + +import java.sql.CallableStatement; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; + +/** + * @author zouzhiwen + */ +@MappedTypes(float[].class) +public class VectorFloatTypeHandler extends BaseTypeHandler { + + @Override + public void setNonNullParameter(PreparedStatement ps, int i, float[] parameter, JdbcType jdbcType) throws SQLException { + PGobject pgObject = new PGobject(); + pgObject.setType("vector"); + pgObject.setValue(arrayToString(parameter)); + ps.setObject(i, pgObject); + } + + @Override + public float[] getNullableResult(ResultSet rs, String columnName) throws SQLException { + String value = rs.getString(columnName); + return parseVector(value); + } + + @Override + public float[] getNullableResult(ResultSet rs, int columnIndex) throws SQLException { + String value = rs.getString(columnIndex); + return parseVector(value); + } + + @Override + public float[] getNullableResult(CallableStatement cs, int columnIndex) throws SQLException { + String value = cs.getString(columnIndex); + return parseVector(value); + } + + private String arrayToString(float[] values) { + StringBuilder builder = new StringBuilder(); + builder.append('['); + for (int i = 0; i < values.length; i++) { + if (i > 0) { + builder.append(','); + } + builder.append(values[i]); + } + builder.append(']'); + return builder.toString(); + } + + private float[] parseVector(String value) { + if (value == null || value.isEmpty()) { + return new float[0]; + } + String content = value.substring(1, value.length() - 1); + if (content.isEmpty()) { + return new float[0]; + } + String[] parts = content.split(","); + float[] result = new float[parts.length]; + for (int i = 0; i < parts.length; i++) { + result[i] = Float.parseFloat(parts[i].trim()); + } + return result; + } +} diff --git a/src/main/resources/application.yml b/src/main/resources/application.yml new file mode 100644 index 0000000..478d801 --- /dev/null +++ b/src/main/resources/application.yml @@ -0,0 +1,28 @@ +server: + port: 8010 + +spring: + config: + import: classpath:domains.yml + datasource: + url: jdbc:postgresql://localhost:5432/assistant + username: postgres + password: postgres + ai: + ollama: + base-url: http://localhost:11434 + chat: + options: + model: qwen2.5:7b + embedding: + options: + model: bge-m3 + +assistant: + default-tenant-id: customer_a + routing: + use-llm: true + fallback-domain: inventory + vector: + table: assistant_vectors + dimension: 1024 diff --git a/src/main/resources/data/sample_inventory_customer_a.jsonl b/src/main/resources/data/sample_inventory_customer_a.jsonl new file mode 100644 index 0000000..e311870 --- /dev/null +++ b/src/main/resources/data/sample_inventory_customer_a.jsonl @@ -0,0 +1,2 @@ +{"owner":"CUST_A","warehouse":"WH01","sku":"SKU-1001","qty":120,"location":"A01-01-01","lot":"LOT-202501","status":"AVAILABLE","updated_at":"2025-01-20 10:00:00"} +{"owner":"CUST_A","warehouse":"WH01","sku":"SKU-2002","qty":45,"location":"B02-03-04","lot":"LOT-202502","status":"HOLD","updated_at":"2025-01-20 11:30:00"} diff --git a/src/main/resources/data/sample_orders_customer_a.jsonl b/src/main/resources/data/sample_orders_customer_a.jsonl new file mode 100644 index 0000000..f1a6421 --- /dev/null +++ b/src/main/resources/data/sample_orders_customer_a.jsonl @@ -0,0 +1,2 @@ +{"order_no":"SO-10001","owner":"CUST_A","warehouse":"WH01","status":"CREATED","qty":12,"created_at":"2025-01-18 09:10:00"} +{"order_no":"SO-10002","owner":"CUST_A","warehouse":"WH01","status":"SHIPPED","qty":5,"created_at":"2025-01-19 16:40:00"} diff --git a/src/main/resources/data/sample_returns_customer_a.jsonl b/src/main/resources/data/sample_returns_customer_a.jsonl new file mode 100644 index 0000000..9bfc879 --- /dev/null +++ b/src/main/resources/data/sample_returns_customer_a.jsonl @@ -0,0 +1,2 @@ +{"return_no":"RT-9001","owner":"CUST_A","warehouse":"WH01","status":"RECEIVED","qty":2,"reason":"damaged","created_at":"2025-01-20 15:25:00"} +{"return_no":"RT-9002","owner":"CUST_A","warehouse":"WH01","status":"PENDING","qty":1,"reason":"wrong_item","created_at":"2025-01-21 10:05:00"} diff --git a/src/main/resources/domains.yml b/src/main/resources/domains.yml new file mode 100644 index 0000000..c89d0f2 --- /dev/null +++ b/src/main/resources/domains.yml @@ -0,0 +1,73 @@ +assistant: + domains: + - id: inventory + name: Inventory + query-mode: live + keywords: ["库存", "库存量", "库位", "库存查询", "qty", "stock"] + fields: ["owner", "warehouse", "sku", "qty", "location", "lot", "status", "updated_at"] + sources: + - type: jsonl + path: classpath:data/sample_inventory_customer_a.jsonl + - type: api + url: http://127.0.0.1:8080/internal/inventory/query + method: POST + headers: + Authorization: Bearer REPLACE_ME + body: + tenant_id: "{tenant_id}" + query: "{question}" + items-path: data.items + - id: orders + name: Orders + query-mode: live + keywords: ["订单", "出库", "发货", "order", "shipment"] + fields: ["order_no", "owner", "warehouse", "status", "qty", "created_at"] + status-mappings: + CREATED: "已创建" + PICKING: "拣货中" + PACKED: "已打包" + SHIPPED: "已发货" + CANCELLED: "已取消" + query-schema: + allowed-params: ["time_range", "status", "owner", "warehouse", "metric"] + enums: + status: ["CREATED", "PICKING", "PACKED", "SHIPPED", "CANCELLED"] + sources: + - type: jsonl + path: classpath:data/sample_orders_customer_a.jsonl + - type: api + url: http://127.0.0.1:8080/internal/orders/query + method: POST + headers: + Authorization: Bearer REPLACE_ME + body: + tenant_id: "{tenant_id}" + query: "{question}" + time_from: "{time_from}" + time_to: "{time_to}" + status: "{status}" + items-path: data.items + - id: returns + name: Returns + query-mode: live + keywords: ["退货", "退单", "return", "refund"] + fields: ["return_no", "owner", "warehouse", "status", "qty", "reason", "created_at"] + query-schema: + allowed-params: ["time_range", "status", "owner", "warehouse", "reason"] + enums: + status: ["PENDING", "RECEIVED", "REJECTED"] + sources: + - type: jsonl + path: classpath:data/sample_returns_customer_a.jsonl + - type: api + url: http://127.0.0.1:8080/internal/returns/query + method: POST + headers: + Authorization: Bearer REPLACE_ME + body: + tenant_id: "{tenant_id}" + query: "{question}" + time_from: "{time_from}" + time_to: "{time_to}" + status: "{status}" + items-path: data.items diff --git a/src/main/resources/mapper/VectorDocumentMapper.xml b/src/main/resources/mapper/VectorDocumentMapper.xml new file mode 100644 index 0000000..bf68703 --- /dev/null +++ b/src/main/resources/mapper/VectorDocumentMapper.xml @@ -0,0 +1,16 @@ + + + + + + +