Commit f8d1fa9b49711da2fef27a6adc7023dede2694db
Committed by
Gitee
!16 文档向量化demo
Merge pull request !16 from 小新/master-jdk21-ai
Showing
8 changed files
with
637 additions
and
0 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/DocService.java
0 → 100644
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/DocServiceImpl.java
0 → 100644
| 1 | +package cn.iocoder.yudao.module.ai.service.knowledge; | |
| 2 | + | |
| 3 | +import jakarta.annotation.Resource; | |
| 4 | +import lombok.extern.slf4j.Slf4j; | |
| 5 | +import org.springframework.ai.document.Document; | |
| 6 | +import org.springframework.ai.reader.tika.TikaDocumentReader; | |
| 7 | +import org.springframework.ai.transformer.splitter.TokenTextSplitter; | |
| 8 | +import org.springframework.ai.vectorstore.RedisVectorStore; | |
| 9 | +import org.springframework.beans.factory.annotation.Value; | |
| 10 | +import org.springframework.stereotype.Service; | |
| 11 | + | |
| 12 | +import java.util.List; | |
| 13 | + | |
| 14 | +/** | |
| 15 | + * AI 知识库 Service 实现类 | |
| 16 | + * | |
| 17 | + * @author xiaoxin | |
| 18 | + */ | |
| 19 | +@Service | |
| 20 | +@Slf4j | |
| 21 | +public class DocServiceImpl implements DocService { | |
| 22 | + | |
| 23 | + @Resource | |
| 24 | + RedisVectorStore vectorStore; | |
| 25 | + @Resource | |
| 26 | + TokenTextSplitter tokenTextSplitter; | |
| 27 | + | |
| 28 | + // TODO @xin 临时测试用,后续删 | |
| 29 | + @Value("classpath:/webapp/test/Fel.pdf") | |
| 30 | + private org.springframework.core.io.Resource data; | |
| 31 | + | |
| 32 | + | |
| 33 | + @Override | |
| 34 | + public void embeddingDoc() { | |
| 35 | + // 读取文件 | |
| 36 | + org.springframework.core.io.Resource file = data; | |
| 37 | + TikaDocumentReader loader = new TikaDocumentReader(file); | |
| 38 | + List<Document> documents = loader.get(); | |
| 39 | + // 文档分段 | |
| 40 | + List<Document> segments = tokenTextSplitter.apply(documents); | |
| 41 | + // 向量化并存储 | |
| 42 | + vectorStore.add(segments); | |
| 43 | + } | |
| 44 | +} | ... | ... |
yudao-module-ai/yudao-spring-boot-starter-ai/pom.xml
| ... | ... | @@ -41,6 +41,28 @@ |
| 41 | 41 | </dependency> |
| 42 | 42 | |
| 43 | 43 | <dependency> |
| 44 | + <groupId>org.springframework.ai</groupId> | |
| 45 | + <artifactId>spring-ai-transformers-spring-boot-starter</artifactId> | |
| 46 | + <version>${spring-ai.version}</version> | |
| 47 | + </dependency> | |
| 48 | + <dependency> | |
| 49 | + <groupId>org.springframework.ai</groupId> | |
| 50 | + <artifactId>spring-ai-tika-document-reader</artifactId> | |
| 51 | + <version>${spring-ai.version}</version> | |
| 52 | + </dependency> | |
| 53 | + <dependency> | |
| 54 | + <groupId>org.springframework.ai</groupId> | |
| 55 | + <artifactId>spring-ai-redis-store</artifactId> | |
| 56 | + <version>${spring-ai.version}</version> | |
| 57 | + </dependency> | |
| 58 | + <dependency> | |
| 59 | + <groupId>org.springframework.data</groupId> | |
| 60 | + <artifactId>spring-data-redis</artifactId> | |
| 61 | + <optional>true</optional> | |
| 62 | + </dependency> | |
| 63 | + | |
| 64 | + | |
| 65 | + <dependency> | |
| 44 | 66 | <groupId>cn.iocoder.boot</groupId> |
| 45 | 67 | <artifactId>yudao-common</artifactId> |
| 46 | 68 | </dependency> | ... | ... |
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/cn/iocoder/yudao/framework/ai/config/YudaoAiAutoConfiguration.java
| ... | ... | @@ -10,11 +10,18 @@ import cn.iocoder.yudao.framework.ai.core.model.xinghuo.XingHuoChatModel; |
| 10 | 10 | import cn.iocoder.yudao.framework.ai.core.model.xinghuo.XingHuoChatOptions; |
| 11 | 11 | import com.alibaba.cloud.ai.tongyi.TongYiAutoConfiguration; |
| 12 | 12 | import lombok.extern.slf4j.Slf4j; |
| 13 | +import org.springframework.ai.autoconfigure.vectorstore.redis.RedisVectorStoreProperties; | |
| 14 | +import org.springframework.ai.document.MetadataMode; | |
| 15 | +import org.springframework.ai.transformer.splitter.TokenTextSplitter; | |
| 16 | +import org.springframework.ai.transformers.TransformersEmbeddingModel; | |
| 17 | +import org.springframework.ai.vectorstore.RedisVectorStore; | |
| 13 | 18 | import org.springframework.boot.autoconfigure.AutoConfiguration; |
| 14 | 19 | import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; |
| 20 | +import org.springframework.boot.autoconfigure.data.redis.RedisProperties; | |
| 15 | 21 | import org.springframework.boot.context.properties.EnableConfigurationProperties; |
| 16 | 22 | import org.springframework.context.annotation.Bean; |
| 17 | 23 | import org.springframework.context.annotation.Import; |
| 24 | +import redis.clients.jedis.JedisPooled; | |
| 18 | 25 | |
| 19 | 26 | /** |
| 20 | 27 | * 芋道 AI 自动配置 |
| ... | ... | @@ -73,4 +80,33 @@ public class YudaoAiAutoConfiguration { |
| 73 | 80 | return new SunoApi(yudaoAiProperties.getSuno().getBaseUrl()); |
| 74 | 81 | } |
| 75 | 82 | |
| 83 | + // ========== rag 相关 ========== | |
| 84 | + @Bean | |
| 85 | + public TransformersEmbeddingModel transformersEmbeddingClient() { | |
| 86 | + return new TransformersEmbeddingModel(MetadataMode.EMBED); | |
| 87 | + } | |
| 88 | + | |
| 89 | + /** | |
| 90 | + * 我们启动有加载很多 Embedding 模型,不晓得取哪个好,先 new 个 TransformersEmbeddingModel 跑 | |
| 91 | + */ | |
| 92 | + @Bean | |
| 93 | + public RedisVectorStore vectorStore(TransformersEmbeddingModel transformersEmbeddingModel, RedisVectorStoreProperties properties, | |
| 94 | + RedisProperties redisProperties) { | |
| 95 | + var config = RedisVectorStore.RedisVectorStoreConfig.builder() | |
| 96 | + .withIndexName(properties.getIndex()) | |
| 97 | + .withPrefix(properties.getPrefix()) | |
| 98 | + .build(); | |
| 99 | + | |
| 100 | + RedisVectorStore redisVectorStore = new RedisVectorStore(config, transformersEmbeddingModel, | |
| 101 | + new JedisPooled(redisProperties.getHost(), redisProperties.getPort()), | |
| 102 | + properties.isInitializeSchema()); | |
| 103 | + redisVectorStore.afterPropertiesSet(); | |
| 104 | + return redisVectorStore; | |
| 105 | + } | |
| 106 | + | |
| 107 | + @Bean | |
| 108 | + public TokenTextSplitter tokenTextSplitter() { | |
| 109 | + return new TokenTextSplitter(500, 100, 5, 10000, true); | |
| 110 | + } | |
| 111 | + | |
| 76 | 112 | } |
| 77 | 113 | \ No newline at end of file | ... | ... |
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/org/springframework/ai/autoconfigure/vectorstore/redis/RedisVectorStoreAutoConfiguration.java
0 → 100644
| 1 | +/* | |
| 2 | + * Copyright 2023 - 2024 the original author or authors. | |
| 3 | + * | |
| 4 | + * Licensed under the Apache License, Version 2.0 (the "License"); | |
| 5 | + * you may not use this file except in compliance with the License. | |
| 6 | + * You may obtain a copy of the License at | |
| 7 | + * | |
| 8 | + * https://www.apache.org/licenses/LICENSE-2.0 | |
| 9 | + * | |
| 10 | + * Unless required by applicable law or agreed to in writing, software | |
| 11 | + * distributed under the License is distributed on an "AS IS" BASIS, | |
| 12 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| 13 | + * See the License for the specific language governing permissions and | |
| 14 | + * limitations under the License. | |
| 15 | + */ | |
| 16 | +package org.springframework.ai.autoconfigure.vectorstore.redis; | |
| 17 | + | |
| 18 | +import org.springframework.ai.embedding.EmbeddingModel; | |
| 19 | +import org.springframework.ai.vectorstore.RedisVectorStore; | |
| 20 | +import org.springframework.ai.vectorstore.RedisVectorStore.RedisVectorStoreConfig; | |
| 21 | +import org.springframework.boot.autoconfigure.AutoConfiguration; | |
| 22 | +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; | |
| 23 | +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; | |
| 24 | +import org.springframework.boot.autoconfigure.data.redis.RedisAutoConfiguration; | |
| 25 | +import org.springframework.boot.context.properties.EnableConfigurationProperties; | |
| 26 | +import org.springframework.context.annotation.Bean; | |
| 27 | +import org.springframework.data.redis.connection.jedis.JedisConnectionFactory; | |
| 28 | +import redis.clients.jedis.JedisPooled; | |
| 29 | + | |
| 30 | +/** | |
| 31 | + * TODO @xin 先拿 spring-ai 最新代码覆盖,1.0.0-M1 跟 redis 自动配置会冲突 | |
| 32 | + * | |
| 33 | + * @author Christian Tzolov | |
| 34 | + * @author Eddú Meléndez | |
| 35 | + */ | |
| 36 | +@AutoConfiguration(after = RedisAutoConfiguration.class) | |
| 37 | +@ConditionalOnClass({JedisPooled.class, JedisConnectionFactory.class, RedisVectorStore.class, EmbeddingModel.class}) | |
| 38 | +//@ConditionalOnBean(JedisConnectionFactory.class) | |
| 39 | +@EnableConfigurationProperties(RedisVectorStoreProperties.class) | |
| 40 | +public class RedisVectorStoreAutoConfiguration { | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + @Bean | |
| 45 | + @ConditionalOnMissingBean | |
| 46 | + public RedisVectorStore vectorStore(EmbeddingModel embeddingModel, RedisVectorStoreProperties properties, | |
| 47 | + JedisConnectionFactory jedisConnectionFactory) { | |
| 48 | + | |
| 49 | + var config = RedisVectorStoreConfig.builder() | |
| 50 | + .withIndexName(properties.getIndex()) | |
| 51 | + .withPrefix(properties.getPrefix()) | |
| 52 | + .build(); | |
| 53 | + | |
| 54 | + return new RedisVectorStore(config, embeddingModel, | |
| 55 | + new JedisPooled(jedisConnectionFactory.getHostName(), jedisConnectionFactory.getPort()), | |
| 56 | + properties.isInitializeSchema()); | |
| 57 | + } | |
| 58 | + | |
| 59 | +} | ... | ... |
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/java/org/springframework/ai/vectorstore/RedisVectorStore.java
0 → 100644
| 1 | +/* | |
| 2 | + * Copyright 2023 - 2024 the original author or authors. | |
| 3 | + * | |
| 4 | + * Licensed under the Apache License, Version 2.0 (the "License"); | |
| 5 | + * you may not use this file except in compliance with the License. | |
| 6 | + * You may obtain a copy of the License at | |
| 7 | + * | |
| 8 | + * https://www.apache.org/licenses/LICENSE-2.0 | |
| 9 | + * | |
| 10 | + * Unless required by applicable law or agreed to in writing, software | |
| 11 | + * distributed under the License is distributed on an "AS IS" BASIS, | |
| 12 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| 13 | + * See the License for the specific language governing permissions and | |
| 14 | + * limitations under the License. | |
| 15 | + */ | |
| 16 | +package org.springframework.ai.vectorstore; | |
| 17 | + | |
| 18 | +import org.slf4j.Logger; | |
| 19 | +import org.slf4j.LoggerFactory; | |
| 20 | +import org.springframework.ai.document.Document; | |
| 21 | +import org.springframework.ai.embedding.EmbeddingModel; | |
| 22 | +import org.springframework.ai.vectorstore.filter.FilterExpressionConverter; | |
| 23 | +import org.springframework.beans.factory.InitializingBean; | |
| 24 | +import org.springframework.util.Assert; | |
| 25 | +import org.springframework.util.CollectionUtils; | |
| 26 | +import redis.clients.jedis.JedisPooled; | |
| 27 | +import redis.clients.jedis.Pipeline; | |
| 28 | +import redis.clients.jedis.json.Path2; | |
| 29 | +import redis.clients.jedis.search.*; | |
| 30 | +import redis.clients.jedis.search.Schema.FieldType; | |
| 31 | +import redis.clients.jedis.search.schemafields.*; | |
| 32 | +import redis.clients.jedis.search.schemafields.VectorField.VectorAlgorithm; | |
| 33 | + | |
| 34 | +import java.text.MessageFormat; | |
| 35 | +import java.util.*; | |
| 36 | +import java.util.function.Function; | |
| 37 | +import java.util.function.Predicate; | |
| 38 | +import java.util.stream.Collectors; | |
| 39 | + | |
| 40 | +/** | |
| 41 | + * The RedisVectorStore is for managing and querying vector data in a Redis database. It | |
| 42 | + * offers functionalities like adding, deleting, and performing similarity searches on | |
| 43 | + * documents. | |
| 44 | + * | |
| 45 | + * The store utilizes RedisJSON and RedisSearch to handle JSON documents and to index and | |
| 46 | + * search vector data. It supports various vector algorithms (e.g., FLAT, HSNW) for | |
| 47 | + * efficient similarity searches. Additionally, it allows for custom metadata fields in | |
| 48 | + * the documents to be stored alongside the vector and content data. | |
| 49 | + * | |
| 50 | + * This class requires a RedisVectorStoreConfig configuration object for initialization, | |
| 51 | + * which includes settings like Redis URI, index name, field names, and vector algorithms. | |
| 52 | + * It also requires an EmbeddingModel to convert documents into embeddings before storing | |
| 53 | + * them. | |
| 54 | + * | |
| 55 | + * @author Julien Ruaux | |
| 56 | + * @author Christian Tzolov | |
| 57 | + * @author Eddú Meléndez | |
| 58 | + * @see VectorStore | |
| 59 | + * @see RedisVectorStoreConfig | |
| 60 | + * @see EmbeddingModel | |
| 61 | + */ | |
| 62 | +public class RedisVectorStore implements VectorStore, InitializingBean { | |
| 63 | + | |
| 64 | + public enum Algorithm { | |
| 65 | + | |
| 66 | + FLAT, HSNW | |
| 67 | + | |
| 68 | + } | |
| 69 | + | |
| 70 | + public record MetadataField(String name, FieldType fieldType) { | |
| 71 | + | |
| 72 | + public static MetadataField text(String name) { | |
| 73 | + return new MetadataField(name, FieldType.TEXT); | |
| 74 | + } | |
| 75 | + | |
| 76 | + public static MetadataField numeric(String name) { | |
| 77 | + return new MetadataField(name, FieldType.NUMERIC); | |
| 78 | + } | |
| 79 | + | |
| 80 | + public static MetadataField tag(String name) { | |
| 81 | + return new MetadataField(name, FieldType.TAG); | |
| 82 | + } | |
| 83 | + | |
| 84 | + } | |
| 85 | + | |
| 86 | + /** | |
| 87 | + * Configuration for the Redis vector store. | |
| 88 | + */ | |
| 89 | + public static final class RedisVectorStoreConfig { | |
| 90 | + | |
| 91 | + private final String indexName; | |
| 92 | + | |
| 93 | + private final String prefix; | |
| 94 | + | |
| 95 | + private final String contentFieldName; | |
| 96 | + | |
| 97 | + private final String embeddingFieldName; | |
| 98 | + | |
| 99 | + private final Algorithm vectorAlgorithm; | |
| 100 | + | |
| 101 | + private final List<MetadataField> metadataFields; | |
| 102 | + | |
| 103 | + private RedisVectorStoreConfig() { | |
| 104 | + this(builder()); | |
| 105 | + } | |
| 106 | + | |
| 107 | + private RedisVectorStoreConfig(Builder builder) { | |
| 108 | + this.indexName = builder.indexName; | |
| 109 | + this.prefix = builder.prefix; | |
| 110 | + this.contentFieldName = builder.contentFieldName; | |
| 111 | + this.embeddingFieldName = builder.embeddingFieldName; | |
| 112 | + this.vectorAlgorithm = builder.vectorAlgorithm; | |
| 113 | + this.metadataFields = builder.metadataFields; | |
| 114 | + } | |
| 115 | + | |
| 116 | + /** | |
| 117 | + * Start building a new configuration. | |
| 118 | + * @return The entry point for creating a new configuration. | |
| 119 | + */ | |
| 120 | + public static Builder builder() { | |
| 121 | + | |
| 122 | + return new Builder(); | |
| 123 | + } | |
| 124 | + | |
| 125 | + /** | |
| 126 | + * {@return the default config} | |
| 127 | + */ | |
| 128 | + public static RedisVectorStoreConfig defaultConfig() { | |
| 129 | + | |
| 130 | + return builder().build(); | |
| 131 | + } | |
| 132 | + | |
| 133 | + public static class Builder { | |
| 134 | + | |
| 135 | + private String indexName = DEFAULT_INDEX_NAME; | |
| 136 | + | |
| 137 | + private String prefix = DEFAULT_PREFIX; | |
| 138 | + | |
| 139 | + private String contentFieldName = DEFAULT_CONTENT_FIELD_NAME; | |
| 140 | + | |
| 141 | + private String embeddingFieldName = DEFAULT_EMBEDDING_FIELD_NAME; | |
| 142 | + | |
| 143 | + private Algorithm vectorAlgorithm = DEFAULT_VECTOR_ALGORITHM; | |
| 144 | + | |
| 145 | + private List<MetadataField> metadataFields = new ArrayList<>(); | |
| 146 | + | |
| 147 | + private Builder() { | |
| 148 | + } | |
| 149 | + | |
| 150 | + /** | |
| 151 | + * Configures the Redis index name to use. | |
| 152 | + * @param name the index name to use | |
| 153 | + * @return this builder | |
| 154 | + */ | |
| 155 | + public Builder withIndexName(String name) { | |
| 156 | + this.indexName = name; | |
| 157 | + return this; | |
| 158 | + } | |
| 159 | + | |
| 160 | + /** | |
| 161 | + * Configures the Redis key prefix to use (default: "embedding:"). | |
| 162 | + * @param prefix the prefix to use | |
| 163 | + * @return this builder | |
| 164 | + */ | |
| 165 | + public Builder withPrefix(String prefix) { | |
| 166 | + this.prefix = prefix; | |
| 167 | + return this; | |
| 168 | + } | |
| 169 | + | |
| 170 | + /** | |
| 171 | + * Configures the Redis content field name to use. | |
| 172 | + * @param name the content field name to use | |
| 173 | + * @return this builder | |
| 174 | + */ | |
| 175 | + public Builder withContentFieldName(String name) { | |
| 176 | + this.contentFieldName = name; | |
| 177 | + return this; | |
| 178 | + } | |
| 179 | + | |
| 180 | + /** | |
| 181 | + * Configures the Redis embedding field name to use. | |
| 182 | + * @param name the embedding field name to use | |
| 183 | + * @return this builder | |
| 184 | + */ | |
| 185 | + public Builder withEmbeddingFieldName(String name) { | |
| 186 | + this.embeddingFieldName = name; | |
| 187 | + return this; | |
| 188 | + } | |
| 189 | + | |
| 190 | + /** | |
| 191 | + * Configures the Redis vector algorithmto use. | |
| 192 | + * @param algorithm the vector algorithm to use | |
| 193 | + * @return this builder | |
| 194 | + */ | |
| 195 | + public Builder withVectorAlgorithm(Algorithm algorithm) { | |
| 196 | + this.vectorAlgorithm = algorithm; | |
| 197 | + return this; | |
| 198 | + } | |
| 199 | + | |
| 200 | + public Builder withMetadataFields(MetadataField... fields) { | |
| 201 | + return withMetadataFields(Arrays.asList(fields)); | |
| 202 | + } | |
| 203 | + | |
| 204 | + public Builder withMetadataFields(List<MetadataField> fields) { | |
| 205 | + this.metadataFields = fields; | |
| 206 | + return this; | |
| 207 | + } | |
| 208 | + | |
| 209 | + /** | |
| 210 | + * {@return the immutable configuration} | |
| 211 | + */ | |
| 212 | + public RedisVectorStoreConfig build() { | |
| 213 | + | |
| 214 | + return new RedisVectorStoreConfig(this); | |
| 215 | + } | |
| 216 | + | |
| 217 | + } | |
| 218 | + | |
| 219 | + } | |
| 220 | + | |
| 221 | + private final boolean initializeSchema; | |
| 222 | + | |
| 223 | + public static final String DEFAULT_INDEX_NAME = "spring-ai-index"; | |
| 224 | + | |
| 225 | + public static final String DEFAULT_CONTENT_FIELD_NAME = "content"; | |
| 226 | + | |
| 227 | + public static final String DEFAULT_EMBEDDING_FIELD_NAME = "embedding"; | |
| 228 | + | |
| 229 | + public static final String DEFAULT_PREFIX = "embedding:"; | |
| 230 | + | |
| 231 | + public static final Algorithm DEFAULT_VECTOR_ALGORITHM = Algorithm.HSNW; | |
| 232 | + | |
| 233 | + private static final String QUERY_FORMAT = "%s=>[KNN %s @%s $%s AS %s]"; | |
| 234 | + | |
| 235 | + private static final Path2 JSON_SET_PATH = Path2.of("$"); | |
| 236 | + | |
| 237 | + private static final String JSON_PATH_PREFIX = "$."; | |
| 238 | + | |
| 239 | + private static final Logger logger = LoggerFactory.getLogger(RedisVectorStore.class); | |
| 240 | + | |
| 241 | + private static final Predicate<Object> RESPONSE_OK = Predicate.isEqual("OK"); | |
| 242 | + | |
| 243 | + private static final Predicate<Object> RESPONSE_DEL_OK = Predicate.isEqual(1l); | |
| 244 | + | |
| 245 | + private static final String VECTOR_TYPE_FLOAT32 = "FLOAT32"; | |
| 246 | + | |
| 247 | + private static final String EMBEDDING_PARAM_NAME = "BLOB"; | |
| 248 | + | |
| 249 | + public static final String DISTANCE_FIELD_NAME = "vector_score"; | |
| 250 | + | |
| 251 | + private static final String DEFAULT_DISTANCE_METRIC = "COSINE"; | |
| 252 | + | |
| 253 | + private final JedisPooled jedis; | |
| 254 | + | |
| 255 | + private final EmbeddingModel embeddingModel; | |
| 256 | + | |
| 257 | + private final RedisVectorStoreConfig config; | |
| 258 | + | |
| 259 | + private FilterExpressionConverter filterExpressionConverter; | |
| 260 | + | |
| 261 | + public RedisVectorStore(RedisVectorStoreConfig config, EmbeddingModel embeddingModel, JedisPooled jedis, | |
| 262 | + boolean initializeSchema) { | |
| 263 | + | |
| 264 | + Assert.notNull(config, "Config must not be null"); | |
| 265 | + Assert.notNull(embeddingModel, "Embedding model must not be null"); | |
| 266 | + this.initializeSchema = initializeSchema; | |
| 267 | + | |
| 268 | + this.jedis = jedis; | |
| 269 | + this.embeddingModel = embeddingModel; | |
| 270 | + this.config = config; | |
| 271 | + this.filterExpressionConverter = new RedisFilterExpressionConverter(this.config.metadataFields); | |
| 272 | + } | |
| 273 | + | |
| 274 | + public JedisPooled getJedis() { | |
| 275 | + return this.jedis; | |
| 276 | + } | |
| 277 | + | |
| 278 | + @Override | |
| 279 | + public void add(List<Document> documents) { | |
| 280 | + try (Pipeline pipeline = this.jedis.pipelined()) { | |
| 281 | + for (Document document : documents) { | |
| 282 | + var embedding = this.embeddingModel.embed(document); | |
| 283 | + document.setEmbedding(embedding); | |
| 284 | + | |
| 285 | + var fields = new HashMap<String, Object>(); | |
| 286 | + fields.put(this.config.embeddingFieldName, embedding); | |
| 287 | + fields.put(this.config.contentFieldName, document.getContent()); | |
| 288 | + fields.putAll(document.getMetadata()); | |
| 289 | + pipeline.jsonSetWithEscape(key(document.getId()), JSON_SET_PATH, fields); | |
| 290 | + } | |
| 291 | + List<Object> responses = pipeline.syncAndReturnAll(); | |
| 292 | + Optional<Object> errResponse = responses.stream().filter(Predicate.not(RESPONSE_OK)).findAny(); | |
| 293 | + if (errResponse.isPresent()) { | |
| 294 | + String message = MessageFormat.format("Could not add document: {0}", errResponse.get()); | |
| 295 | + if (logger.isErrorEnabled()) { | |
| 296 | + logger.error(message); | |
| 297 | + } | |
| 298 | + throw new RuntimeException(message); | |
| 299 | + } | |
| 300 | + } | |
| 301 | + } | |
| 302 | + | |
| 303 | + private String key(String id) { | |
| 304 | + return this.config.prefix + id; | |
| 305 | + } | |
| 306 | + | |
| 307 | + @Override | |
| 308 | + public Optional<Boolean> delete(List<String> idList) { | |
| 309 | + try (Pipeline pipeline = this.jedis.pipelined()) { | |
| 310 | + for (String id : idList) { | |
| 311 | + pipeline.jsonDel(key(id)); | |
| 312 | + } | |
| 313 | + List<Object> responses = pipeline.syncAndReturnAll(); | |
| 314 | + Optional<Object> errResponse = responses.stream().filter(Predicate.not(RESPONSE_DEL_OK)).findAny(); | |
| 315 | + if (errResponse.isPresent()) { | |
| 316 | + if (logger.isErrorEnabled()) { | |
| 317 | + logger.error("Could not delete document: {}", errResponse.get()); | |
| 318 | + } | |
| 319 | + return Optional.of(false); | |
| 320 | + } | |
| 321 | + return Optional.of(true); | |
| 322 | + } | |
| 323 | + } | |
| 324 | + | |
| 325 | + @Override | |
| 326 | + public List<Document> similaritySearch(SearchRequest request) { | |
| 327 | + | |
| 328 | + Assert.isTrue(request.getTopK() > 0, "The number of documents to returned must be greater than zero"); | |
| 329 | + Assert.isTrue(request.getSimilarityThreshold() >= 0 && request.getSimilarityThreshold() <= 1, | |
| 330 | + "The similarity score is bounded between 0 and 1; least to most similar respectively."); | |
| 331 | + | |
| 332 | + String filter = nativeExpressionFilter(request); | |
| 333 | + | |
| 334 | + String queryString = String.format(QUERY_FORMAT, filter, request.getTopK(), this.config.embeddingFieldName, | |
| 335 | + EMBEDDING_PARAM_NAME, DISTANCE_FIELD_NAME); | |
| 336 | + | |
| 337 | + List<String> returnFields = new ArrayList<>(); | |
| 338 | + this.config.metadataFields.stream().map(MetadataField::name).forEach(returnFields::add); | |
| 339 | + returnFields.add(this.config.embeddingFieldName); | |
| 340 | + returnFields.add(this.config.contentFieldName); | |
| 341 | + returnFields.add(DISTANCE_FIELD_NAME); | |
| 342 | + var embedding = toFloatArray(this.embeddingModel.embed(request.getQuery())); | |
| 343 | + Query query = new Query(queryString).addParam(EMBEDDING_PARAM_NAME, RediSearchUtil.toByteArray(embedding)) | |
| 344 | + .returnFields(returnFields.toArray(new String[0])) | |
| 345 | + .setSortBy(DISTANCE_FIELD_NAME, true) | |
| 346 | + .dialect(2); | |
| 347 | + | |
| 348 | + SearchResult result = this.jedis.ftSearch(this.config.indexName, query); | |
| 349 | + return result.getDocuments() | |
| 350 | + .stream() | |
| 351 | + .filter(d -> similarityScore(d) >= request.getSimilarityThreshold()) | |
| 352 | + .map(this::toDocument) | |
| 353 | + .toList(); | |
| 354 | + } | |
| 355 | + | |
| 356 | + private Document toDocument(redis.clients.jedis.search.Document doc) { | |
| 357 | + var id = doc.getId().substring(this.config.prefix.length()); | |
| 358 | + var content = doc.hasProperty(this.config.contentFieldName) ? doc.getString(this.config.contentFieldName) | |
| 359 | + : null; | |
| 360 | + Map<String, Object> metadata = this.config.metadataFields.stream() | |
| 361 | + .map(MetadataField::name) | |
| 362 | + .filter(doc::hasProperty) | |
| 363 | + .collect(Collectors.toMap(Function.identity(), doc::getString)); | |
| 364 | + metadata.put(DISTANCE_FIELD_NAME, 1 - similarityScore(doc)); | |
| 365 | + return new Document(id, content, metadata); | |
| 366 | + } | |
| 367 | + | |
| 368 | + private float similarityScore(redis.clients.jedis.search.Document doc) { | |
| 369 | + return (2 - Float.parseFloat(doc.getString(DISTANCE_FIELD_NAME))) / 2; | |
| 370 | + } | |
| 371 | + | |
| 372 | + private String nativeExpressionFilter(SearchRequest request) { | |
| 373 | + if (request.getFilterExpression() == null) { | |
| 374 | + return "*"; | |
| 375 | + } | |
| 376 | + return "(" + this.filterExpressionConverter.convertExpression(request.getFilterExpression()) + ")"; | |
| 377 | + } | |
| 378 | + | |
| 379 | + @Override | |
| 380 | + public void afterPropertiesSet() { | |
| 381 | + | |
| 382 | + if (!this.initializeSchema) { | |
| 383 | + return; | |
| 384 | + } | |
| 385 | + | |
| 386 | + // If index already exists don't do anything | |
| 387 | + if (this.jedis.ftList().contains(this.config.indexName)) { | |
| 388 | + return; | |
| 389 | + } | |
| 390 | + | |
| 391 | + String response = this.jedis.ftCreate(this.config.indexName, | |
| 392 | + FTCreateParams.createParams().on(IndexDataType.JSON).addPrefix(this.config.prefix), schemaFields()); | |
| 393 | + if (!RESPONSE_OK.test(response)) { | |
| 394 | + String message = MessageFormat.format("Could not create index: {0}", response); | |
| 395 | + throw new RuntimeException(message); | |
| 396 | + } | |
| 397 | + } | |
| 398 | + | |
| 399 | + private Iterable<SchemaField> schemaFields() { | |
| 400 | + Map<String, Object> vectorAttrs = new HashMap<>(); | |
| 401 | + vectorAttrs.put("DIM", this.embeddingModel.dimensions()); | |
| 402 | + vectorAttrs.put("DISTANCE_METRIC", DEFAULT_DISTANCE_METRIC); | |
| 403 | + vectorAttrs.put("TYPE", VECTOR_TYPE_FLOAT32); | |
| 404 | + List<SchemaField> fields = new ArrayList<>(); | |
| 405 | + fields.add(TextField.of(jsonPath(this.config.contentFieldName)).as(this.config.contentFieldName).weight(1.0)); | |
| 406 | + fields.add(VectorField.builder() | |
| 407 | + .fieldName(jsonPath(this.config.embeddingFieldName)) | |
| 408 | + .algorithm(vectorAlgorithm()) | |
| 409 | + .attributes(vectorAttrs) | |
| 410 | + .as(this.config.embeddingFieldName) | |
| 411 | + .build()); | |
| 412 | + | |
| 413 | + if (!CollectionUtils.isEmpty(this.config.metadataFields)) { | |
| 414 | + for (MetadataField field : this.config.metadataFields) { | |
| 415 | + fields.add(schemaField(field)); | |
| 416 | + } | |
| 417 | + } | |
| 418 | + return fields; | |
| 419 | + } | |
| 420 | + | |
| 421 | + private SchemaField schemaField(MetadataField field) { | |
| 422 | + String fieldName = jsonPath(field.name); | |
| 423 | + switch (field.fieldType) { | |
| 424 | + case NUMERIC: | |
| 425 | + return NumericField.of(fieldName).as(field.name); | |
| 426 | + case TAG: | |
| 427 | + return TagField.of(fieldName).as(field.name); | |
| 428 | + case TEXT: | |
| 429 | + return TextField.of(fieldName).as(field.name); | |
| 430 | + default: | |
| 431 | + throw new IllegalArgumentException( | |
| 432 | + MessageFormat.format("Field {0} has unsupported type {1}", field.name, field.fieldType)); | |
| 433 | + } | |
| 434 | + } | |
| 435 | + | |
| 436 | + private VectorAlgorithm vectorAlgorithm() { | |
| 437 | + if (config.vectorAlgorithm == Algorithm.HSNW) { | |
| 438 | + return VectorAlgorithm.HNSW; | |
| 439 | + } | |
| 440 | + return VectorAlgorithm.FLAT; | |
| 441 | + } | |
| 442 | + | |
| 443 | + private String jsonPath(String field) { | |
| 444 | + return JSON_PATH_PREFIX + field; | |
| 445 | + } | |
| 446 | + | |
| 447 | + private static float[] toFloatArray(List<Double> embeddingDouble) { | |
| 448 | + float[] embeddingFloat = new float[embeddingDouble.size()]; | |
| 449 | + int i = 0; | |
| 450 | + for (Double d : embeddingDouble) { | |
| 451 | + embeddingFloat[i++] = d.floatValue(); | |
| 452 | + } | |
| 453 | + return embeddingFloat; | |
| 454 | + } | |
| 455 | + | |
| 456 | +} | |
| 0 | 457 | \ No newline at end of file | ... | ... |
yudao-module-ai/yudao-spring-boot-starter-ai/src/main/resources/webapp/test/Fel.pdf
0 → 100755
No preview for this file type
yudao-server/src/main/resources/application.yaml