diff --git a/java/connectors/semantickernel-connectors-memory-redis/pom.xml b/java/connectors/semantickernel-connectors-memory-redis/pom.xml new file mode 100644 index 0000000000..eb419a16ce --- /dev/null +++ b/java/connectors/semantickernel-connectors-memory-redis/pom.xml @@ -0,0 +1,74 @@ + + + 4.0.0 + + com.microsoft.semantic-kernel + semantickernel-parent + 0.2.10-alpha-SNAPSHOT + ../../pom.xml + + + semantickernel-connectors-memory-redis + + Semantic Kernel Redis Memory Connector + Defines connector for Redis to be used with the Semantic Kernel + + + UTF-8 + UTF-8 + 17 + 17 + 17 + 17 + + + + + com.fasterxml.jackson.core + jackson-core + compile + + + com.fasterxml.jackson.core + jackson-databind + compile + + + com.microsoft.semantic-kernel + semantickernel-api + + + org.mockito + mockito-core + + + org.junit.jupiter + junit-jupiter + + + redis.clients + jedis + 5.1.0 + + + com.microsoft.semantic-kernel + semantickernel-connectors-memory-jdbc + ${project.version} + compile + + + + + + org.apache.maven.plugins + maven-compiler-plugin + + + 17 + 17 + 17 + + + + + \ No newline at end of file diff --git a/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/JsonMemoryRecord.java b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/JsonMemoryRecord.java new file mode 100644 index 0000000000..6a2bf1a63c --- /dev/null +++ b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/JsonMemoryRecord.java @@ -0,0 +1,173 @@ +// Copyright (c) Microsoft. All rights reserved. +package com.microsoft.semantickernel.connectors.memory.redis; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.microsoft.semantickernel.ai.embeddings.Embedding; +import com.microsoft.semantickernel.memory.MemoryRecord; +import com.microsoft.semantickernel.memory.MemoryRecordMetadata; + +import javax.annotation.Nullable; +import java.nio.charset.StandardCharsets; +import java.util.*; + +/** Represents an entry in the Semantic Kernel Memory Table. */ +public class JsonMemoryRecord { + + static final String ID = "Id"; + static final String TEXT = "Text"; + static final String DESCRIPTION = "Description"; + static final String ADDITIONAL_METADATA = "AdditionalMetadata"; + static final String EMBEDDING = "Embedding"; + static final String EXTERNAL_SOURCE_NAME = "ExternalSourceName"; + static final String IS_REFERENCE = "Reference"; + static final String TIME = "Timestamp"; + + private final String id; + private final String text; + private final String description; + private final String additionalMetadata; + private final List embedding; + private final String externalSourceName; + private final boolean isReference; + + @JsonCreator + public JsonMemoryRecord( + @JsonProperty(ID) String id, + @JsonProperty(TEXT) String text, + @JsonProperty(DESCRIPTION) String description, + @JsonProperty(ADDITIONAL_METADATA) String additionalMetadata, + @JsonProperty(EMBEDDING) List embedding, + @JsonProperty(EXTERNAL_SOURCE_NAME) String externalSourceName, + @JsonProperty(IS_REFERENCE) boolean isReference) { + this.id = id; + this.text = text; + this.description = description; + this.additionalMetadata = additionalMetadata; + this.embedding = embedding != null ? embedding : Collections.emptyList(); + this.externalSourceName = externalSourceName; + this.isReference = isReference; + } + + + /** + * Record ID. The record is not filterable to save quota, also SK uses only semantic search. + * + * @return Record ID. + */ + public String getId() { + return id; + } + + /** + * Content is stored here. + * + * @return Content is stored here, {@code null} if not set. + */ + public String getText() { + return text; + } + + /** + * Optional description of the content, e.g. a title. This can be useful when indexing external + * data without pulling in the entire content. + * + * @return Optional description of the content, {@code null} if not set. + */ + public String getDescription() { + return description; + } + + /** + * Additional metadata. Currently, this is a String where you could store serialized data as + * JSON. In future the design might change to allow storing named values and leverage filters. + * + * @return Additional metadata, {@code null} if not set. + */ + public String getAdditionalMetadata() { + return additionalMetadata; + } + + /** + * Embedding vector. + * + * @return Embedding vector. + */ + public List getEmbedding() { + return embedding; + } + + /** + * Name of the external source, in cases where the content and the Id are referenced to external + * information. + * + * @return Name of the external source, in cases where the content and the Id are referenced to + * external information, {@code null} if not set. + */ + public String getExternalSourceName() { + return externalSourceName; + } + + /** + * Whether the record references external information. + * + * @return {@code true} if the record references external information, {@code false} otherwise. + */ + public boolean isReference() { + return isReference; + } + + public MemoryRecordMetadata toMemoryRecordMetadata() { + return new MemoryRecordMetadata( + this.isReference, + decodeId(this.id), + this.text, + this.description, + this.externalSourceName, + this.additionalMetadata); + } + + public static JsonMemoryRecord fromMemoryRecord(MemoryRecord record) { + return new JsonMemoryRecord( + record.getMetadata().getId(), + record.getMetadata().getText(), + record.getMetadata().getDescription(), + record.getMetadata().getAdditionalMetadata(), + record.getEmbedding().getVector(), + record.getMetadata().getExternalSourceName(), + record.getMetadata().isReference()); + } + + public MemoryRecord toMemoryRecord() { + return toMemoryRecord(true); + } + + public MemoryRecord toMemoryRecord(boolean withEmbeddings) { + return new MemoryRecord( + this.toMemoryRecordMetadata(), + new Embedding(withEmbeddings ? this.embedding : Collections.emptyList()), + this.id, + null); + } + + // ACS keys can contain only letters, digits, underscore, dash, equal sign, recommending + // to encode values with a URL-safe algorithm. + // Original Id + // Encoded id + static String encodeId(@Nullable String realId) { + if (realId == null) { + return ""; + } + byte[] bytes = Base64.getUrlEncoder().encode(realId.getBytes(StandardCharsets.UTF_8)); + return new String(bytes, StandardCharsets.UTF_8); + } + + static String decodeId(@Nullable String encodedId) { + if (encodedId == null) { + return ""; + } + byte[] bytes = Base64.getUrlDecoder().decode(encodedId.getBytes(StandardCharsets.UTF_8)); + return new String(bytes, StandardCharsets.UTF_8); + } + +} diff --git a/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisException.java b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisException.java new file mode 100644 index 0000000000..6bd11309cb --- /dev/null +++ b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisException.java @@ -0,0 +1,70 @@ +// Copyright (c) Microsoft. All rights reserved. +package com.microsoft.semantickernel.connectors.memory.redis; + +import com.microsoft.semantickernel.SKException; + +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + +/** Exception thrown by the SQL connector. */ +public class RedisException extends SKException { + + private final ErrorCodes errorCode; + + /** + * Create an exception with a message + * + * @param errorCode The error code + */ + public RedisException(@Nonnull ErrorCodes errorCode) { + + this(errorCode, null, null); + } + + /** + * Create an exception with a message + * + * @param errorCode The error code + * @param message a description of the cause of the exception + */ + public RedisException(@Nonnull ErrorCodes errorCode, @Nullable String message) { + + this(errorCode, message, null); + } + + /** + * Create an exception with a message and a cause + * + * @param errorCode the error code + * @param message a description of the cause of the exception + * @param cause the cause of the exception + */ + public RedisException( + @Nonnull ErrorCodes errorCode, @Nullable String message, @Nullable Throwable cause) { + super(message, cause); + this.errorCode = errorCode; + } + + public ErrorCodes getErrorCode() { + return errorCode; + } + + public enum ErrorCodes { + REDIS_ERROR("Redis error"), + UNKNOWN_ERROR("Unknown error"), + INVALID_EMBEDDING_SIZE("Invalid embedding size"), + INVALID_INDEX_NAME("Invalid index name"), + READ_FAILURE("Read failure"), + WRITE_FAILURE("Write failure"); + + final String message; + + ErrorCodes(String message) { + this.message = message; + } + + public String getMessage() { + return message; + } + } +} diff --git a/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisIndexSchemaParams.java b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisIndexSchemaParams.java new file mode 100644 index 0000000000..1477677fa3 --- /dev/null +++ b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisIndexSchemaParams.java @@ -0,0 +1,8 @@ +package com.microsoft.semantickernel.connectors.memory.redis; + +public class RedisIndexSchemaParams { + + public static final String TYPE = "TYPE"; + public static final String DIM = "DIM"; + public static final String DIST = "DISTANCE_METRIC"; +} diff --git a/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisMemoryRecord.java b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisMemoryRecord.java new file mode 100644 index 0000000000..d0f20a7ef8 --- /dev/null +++ b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisMemoryRecord.java @@ -0,0 +1,97 @@ +// Copyright (c) Microsoft. All rights reserved. +package com.microsoft.semantickernel.connectors.memory.redis; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.microsoft.semantickernel.ai.embeddings.Embedding; +import com.microsoft.semantickernel.memory.DataEntryBase; +import com.microsoft.semantickernel.memory.MemoryRecord; +import redis.clients.jedis.search.Document; + +import java.time.ZonedDateTime; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; + +/** Represents an entry in the Semantic Kernel Memory Table. */ +public class RedisMemoryRecord extends JsonMemoryRecord { + + public RedisMemoryRecord(String id, String text, String description, String additionalMetadata, + List embedding, String externalSourceName, boolean isReference) { + super(id, text, description, additionalMetadata, embedding, externalSourceName, isReference); + } + + /** + * Converts a RedisSearch Document to a HashMap. + * + * @param document A Redis Search Document + * + * @return A HashMap containing the iterator's entries + */ + protected Map documentToMap(Document document) { + return StreamSupport.stream(document.getProperties().spliterator(), false) + .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); + } + + /** + * Gets the embedding information associated with the entry. + * + * @param map A hashmap with keys representing a record + * @param withEmbedding A boolean flag indicating that the embedding should be returned with the metadata + * + * @return A MemoryRecord object. + */ + public static MemoryRecord mapToRecord(Map map, boolean withEmbedding) { + Embedding embedding = Embedding.empty(); + ZonedDateTime time = null; + ObjectMapper mapper = new ObjectMapper(); + try { + if (map.get(TIME) != null) { + time = mapper.readValue(map.getOrDefault(TIME, ZonedDateTime.now().toString()), ZonedDateTime.class); + } + if (withEmbedding) { + List embeddings = mapper.readValue(map.get(EMBEDDING), List.class); + embedding = new Embedding(embeddings); + } + } catch (JsonProcessingException e) { + throw new RedisException( + RedisException.ErrorCodes.REDIS_ERROR, + "Error deserializing Redis entry", + e); + } + + return MemoryRecord.localRecord( + map.get(ID), + map.get(TEXT), + map.get(DESCRIPTION), + embedding, + map.get(ADDITIONAL_METADATA), + map.get(EXTERNAL_SOURCE_NAME), + time + ); + } + + /** + * Gets the embedding information associated with the entry. + * + * @param record A MemoryRecord object + * + * @return A hashmap containing the entries of the record. + */ + public static Map recordToMap(MemoryRecord record) { + + Map map = new HashMap(); + + map.put(ID, record.getMetadata().getId()); + map.put(EXTERNAL_SOURCE_NAME, record.getMetadata().getExternalSourceName()); + map.put(TEXT, record.getMetadata().getText()); + map.put(DESCRIPTION, record.getMetadata().getDescription()); + map.put(EMBEDDING, record.getEmbedding().getVector()); + map.put(ADDITIONAL_METADATA, record.getMetadata().getAdditionalMetadata()); + map.put(IS_REFERENCE, record.getMetadata().isReference()); + + return map; + } +} diff --git a/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisMemoryStore.java b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisMemoryStore.java new file mode 100644 index 0000000000..368e0b9e45 --- /dev/null +++ b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisMemoryStore.java @@ -0,0 +1,572 @@ +package com.microsoft.semantickernel.connectors.memory.redis; + +import com.microsoft.semantickernel.ai.embeddings.Embedding; +import com.microsoft.semantickernel.memory.MemoryException; +import com.microsoft.semantickernel.memory.MemoryException.ErrorCodes; +import com.microsoft.semantickernel.memory.MemoryRecord; +import com.microsoft.semantickernel.memory.MemoryStore; +import reactor.core.publisher.Mono; +import reactor.util.function.Tuple2; +import reactor.util.function.Tuples; +import redis.clients.jedis.JedisPooled; +import redis.clients.jedis.exceptions.JedisDataException; +import redis.clients.jedis.resps.ScanResult; +import redis.clients.jedis.search.*; +import redis.clients.jedis.search.schemafields.VectorField.VectorAlgorithm; +import javax.annotation.Nonnull; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.security.InvalidParameterException; +import java.text.MessageFormat; +import java.time.Instant; +import java.time.OffsetDateTime; +import java.time.ZoneId; +import java.util.*; + +import java.util.function.Function; +import java.util.stream.Collectors; + +/** + * Semantic Memory implementation using Redis Vector Search. For more information about Redis + * Vector Search {@see https://redis.com/solutions/use-cases/vector-database/} + */ +public class RedisMemoryStore implements MemoryStore { + + private static final VectorAlgorithm DefaultIndexAlgorithm = VectorAlgorithm.HNSW; // FLAT or HNSW + private static final String DefaultVectorType = RedisVectorType.FLOAT32; // prefer better accuracy + private static final String DefaultDistanceMetric = RedisVectorDistanceMetric.L2; + private static final Integer DefaultQueryDialect = 2; // this has to be >= 2 for vector functionality + private static final Integer DefaultVectorSize = 1536; // framework model dependent - ada-002 + private static final String SUFFIX = "sk"; + private static final String INDEX = "idx"; + private static final String OK = "OK"; + + private final JedisPooled client; + private Integer vectorSize = DefaultVectorSize; + private VectorAlgorithm vectorIndexAlgorithm = DefaultIndexAlgorithm; + private String vectorDistanceMetric = DefaultDistanceMetric; + private String vectorType = DefaultVectorType; + private Integer queryDialect; + + public static void isXinRange(Integer x, Integer lower, Integer upper, String message) { + if (lower > x || x > upper) { + throw new IllegalArgumentException(MessageFormat.format(message, x)); + } + } + + /** + * Create a new instance of semantic memory using Redis. + * + * @param client A Redis Database client connection. + * @param vectorSize Embedding vector size, defaults to 1536 - framework model dependent - ada-002. + * @param vectorIndexAlgorithm Indexing algorithm for vectors, defaults to "HNSW" + * @param vectorDistanceMetric Metric for measuring vector distances, defaults to "COSINE" + * @param queryDialect Query dialect, must be 2 or greater for vector similarity searching, defaults to 2 + */ + public RedisMemoryStore( + JedisPooled client, + Integer vectorSize, + VectorAlgorithm vectorIndexAlgorithm, + String vectorDistanceMetric, + Integer queryDialect + ) { + + isXinRange(vectorSize, 0, 8192, + "Invalid vector size: {x}. Vector size must be in the range 0-8192."); + isXinRange(queryDialect, 2, 4, + "Invalid query dialect: {x}. Query dialect must be in the range 2-4."); + + this.client = client; + + this.vectorSize = vectorSize; + this.vectorIndexAlgorithm = vectorIndexAlgorithm; + this.vectorDistanceMetric = vectorDistanceMetric; + this.queryDialect = queryDialect; + } + + /** + * Create a new instance of semantic memory using Redis. + * + * @param connectionString Provide connection URL to a Redis instance. + * @param vectorSize Embedding vector size, defaults to 1536. + * @param vectorIndexAlgorithm Indexing algorithm for vectors, defaults to "HNSW" + * @param vectorDistanceMetric Metric for measuring vector distances, defaults to "COSINE" + * @param queryDialect Query dialect, must be 2 or greater for vector similarity searching, defaults to 2 + */ + public RedisMemoryStore( + String connectionString, + Integer vectorSize, + VectorAlgorithm vectorIndexAlgorithm, + String vectorDistanceMetric, + Integer queryDialect) { + + isXinRange(vectorSize, 0, 8192, + "Invalid vector size: {x}. Vector size must be in the range 0-8192."); + isXinRange(queryDialect, 2, 4, + "Invalid query dialect: {x}. Query dialect must be in the range 2-4."); + + this.client = new JedisPooled(connectionString); + + this.vectorSize = vectorSize; + this.vectorIndexAlgorithm = vectorIndexAlgorithm; + this.vectorDistanceMetric = vectorDistanceMetric; + this.queryDialect = queryDialect; + } + + /** + * Returns a Redis key + * + * @param collectionName The name associated with a collection of embeddings. + * @param key The unique id associated with the memory record to get. + * @return A Redis key that identifies a particular Hash of Json object. + */ + private static String getRedisKey(String collectionName, String key) { + + String name = collectionName.toLowerCase(Locale.ROOT); + String id = JsonMemoryRecord.encodeId(key); + return String.format("%s:%s", name, id); + } + + /** + * Returns a long timestamp given a OffsetDateTime + * + * @param timestamp An OffsetDateTime. + * @return A long timestamp or -1 in the case of an error. + */ + private static long toTimestampLong(OffsetDateTime timestamp) { + if (timestamp != null) { + return timestamp.toInstant().toEpochMilli(); + } + return -1; + } + + /** + * Returns an OffsetDateTime given a valid timestamp + * + * @param timestamp A long timestamp value. + * @return An OffsetDateTime if the timestamp is valid, else null. + */ + private static OffsetDateTime toDateTime(long timestamp) { + if (timestamp > 0) { + Instant instant = Instant.ofEpochMilli(timestamp); + return OffsetDateTime.ofInstant(instant, ZoneId.systemDefault()); + } + return null; + } + + @Override + public Mono> getCollectionsAsync() { + return Mono.just(getIndexesAsync().block().stream() + .map(name -> name.substring(0, name.length() - "-sk-idx".length())) + .collect(Collectors.toList())); + } + + /** + * Returns a List of Index(Collection) Names + * + * @return A list of names of the embedding indexes, including suffix, in the database + */ + private Mono> getIndexesAsync() { + return Mono.just(this.client.ftList().stream() + .filter(name -> name.endsWith("-sk-idx")) + .collect(Collectors.toList())); + } + + @Override + public Mono createCollectionAsync(@Nonnull String collectionName) { + // Indexes are created when sending a record; creation requires the size of the embedding vector + return Mono.empty(); + } + + public Mono doesIndexExistAsync(@Nonnull String collectionName) { + + try { + Map info = this.client.ftInfo(collectionName); + return Mono.just(info != null && ! info.isEmpty()); + } catch(Exception e) { + if (!(e instanceof JedisDataException)) { + throw e; + } + return Mono.just(false); + } + // return Mono.just(this.client.ftInfo(createIndexName(collectionName)) != null); + } + + @Override + public Mono doesCollectionExistAsync(@Nonnull String collectionName) { + + Objects.requireNonNull(collectionName); + String normalizedIndexName = createIndexName(collectionName); + return getIndexesAsync() + .map( + list -> + list.stream() + .anyMatch( + name -> + name.equalsIgnoreCase(collectionName) + || name.equalsIgnoreCase( + normalizedIndexName))); + } + + @Override + public Mono deleteCollectionAsync(@Nonnull String collectionName) { + Objects.requireNonNull(collectionName); + try { + this.client.ftDropIndex(createIndexName(collectionName)); + return Mono.empty(); + } catch (Exception e) { + return Mono.error(e); + } + } + + @Override + public Mono getAsync(@Nonnull String collectionName, @Nonnull String key, + boolean withEmbedding) { + + Map map = this.client.hgetAll(String.format("%s:%s", + collectionName, key).toLowerCase(Locale.ROOT)); + if (map == null || map.isEmpty()) { + return Mono.empty(); + } + MemoryRecord record = RedisMemoryRecord.mapToRecord(map, withEmbedding); + return Mono.just(record); + } + + public Mono getInternalAsync(@Nonnull String collectionName, @Nonnull String key, + boolean withEmbedding) { + + Objects.requireNonNull(collectionName); + Objects.requireNonNull(key); + + Map entry = this.client.hgetAll(String.format("%s:%s", collectionName, key)); + + if (entry == null) + return Mono.empty(); + + return Mono.just(RedisMemoryRecord.mapToRecord(entry, withEmbedding)); + } + + /** + * Converts a RedisSearch Document to a HashMap. + * + * @param document A Redis Search Document + * + * @return A HashMap containing the iterator's entries + */ + protected Map documentToMap(Document document) { + + return iterableToMap(document.getProperties()); + } + + /** + * Converts an Iterable> to a HashMap. + * + * @param entryIterator An Iterable> + * + * @return A HashMap containing the iterator's entries + */ + protected Map iterableToMap(Iterable> entryIterator) { + + Map entryMap = new HashMap<>(); + entryIterator.forEach((e) -> entryMap.put(e.getKey(), e.getValue().toString())); + + return entryMap; + } + + protected MemoryRecord documentToMemoryRecord(Document document, boolean withEmbedding) { + + Map map = documentToMap(document); + return RedisMemoryRecord.mapToRecord(map, withEmbedding); + } + + @Override + public Mono> getBatchAsync(@Nonnull String collectionName, + @Nonnull Collection keys, boolean withEmbeddings) { + Objects.requireNonNull(collectionName); + // ACS issues one query per key; redis has several possible calls, fastest will be hkeys followed by hvals + Collection records = keys.stream().map(key -> getAsync(collectionName, key, withEmbeddings).block()) + .collect(Collectors.toList()); + + return Mono.just(records); + } + + @Override + public Mono upsertAsync(@Nonnull String collectionName, @Nonnull MemoryRecord record) { + return upsertRecordAsync(collectionName, record); + } + + /** + * Returns the key of the upserted object (wraps the record in a list and calls the batch function). + * + * @param collectionName The name associated with a collection of embeddings. + * @param record The memory record to upsert. + * @return The key of the newly created record. + */ + public Mono upsertRecordAsync(@Nonnull String collectionName, @Nonnull MemoryRecord record) { + return upsertBatchAsync(collectionName, Collections.singletonList(record)) + .map(Collection::iterator) + .map(Iterator::next); + } + + @Override + public Mono> upsertBatchAsync(@Nonnull String collectionName, + @Nonnull Collection records) { + + Objects.requireNonNull(collectionName); + Objects.requireNonNull(records); + + if (records.isEmpty()) { + return Mono.just(Collections.emptyList()); + } + + doesCollectionExistAsync(collectionName) + .map( + exists -> { + if (! exists) { + int embeddingSize = records.stream() + .map(record -> record.getEmbedding().getVector()) + .map(List::size) + .max(Integer::compareTo) + .orElse(0); + createIndexAsync(collectionName, embeddingSize); + } + return true; + } + ); + + return Mono.just(records.stream().map(record -> { + Map map = RedisMemoryRecord.recordToMap(record); + // we have to change the embeddings to bytes - it's easiest to modify the map + List embedding = record.getEmbedding().getVector(); + map.replace(RedisMemoryRecord.EMBEDDING, embedding, embeddingToBytes(embedding)); + String key = getRedisKey(collectionName, record.getMetadata().getId()); + this.client.hsetObject(key, map); + return key; + }).collect(Collectors.toList())); + } + + @Override + public Mono removeAsync(@Nonnull String collectionName, @Nonnull String key) { + return Mono.fromRunnable( + () -> this.client.hdel(String.format("%s:%s", collectionName, key))); + } + + @Override + public Mono removeBatchAsync(@Nonnull String collectionName, + @Nonnull Collection keys) { + return Mono.fromRunnable( + () -> { + // ScanResult results = this.client.scan(collectionName); + // results.getResult().forEach(key -> this.client.hdel(key)); + keys.forEach(key -> this.client.hdel(String.format("%s:%s", collectionName, key))); + }); + } + + /** + * Create a new search index. + * + * @param collectionName Index name + * @param embeddingSize Size of the embedding vector + * @return A Mono that completes when the index is created + */ + private Mono createIndexAsync( + @Nonnull String collectionName, int embeddingSize) { + if (embeddingSize < 1) { + throw new RedisException( + RedisException.ErrorCodes.INVALID_EMBEDDING_SIZE, "the value must be greater than zero"); + } + + Map attributes = Map.ofEntries( + new AbstractMap.SimpleEntry(RedisIndexSchemaParams.TYPE, vectorType), + new AbstractMap.SimpleEntry(RedisIndexSchemaParams.DIM, embeddingSize), + new AbstractMap.SimpleEntry(RedisIndexSchemaParams.DIST, vectorDistanceMetric)); + + Schema schema = new Schema() + .addVectorField("Embedding", Schema.VectorField.VectorAlgo.FLAT, attributes).as("Embedding") + .addTextField("Id", 1.0).as("Id") + .addTextField("Text", 1.0).as("Text") + .addTextField("Description", 1.0).as("Description") + .addTextField("AdditionalMetadata", 1.0).as("AdditionalMetadata") + .addTextField("ExternalSourceName", 1.0).as("ExternalSourceName"); + + IndexDefinition rule = new IndexDefinition(IndexDefinition.Type.HASH) + .setPrefixes(collectionName.toLowerCase() + ":"); + // customarily redis indexes are named -idx; we're going to suffix them (-sk-idx) so we can find them + String normalizedIndexName = createIndexName(collectionName); + String result = this.client.ftCreate(normalizedIndexName, + IndexOptions.defaultOptions().setDefinition(rule), schema); + + if (result.equals(OK)) { + return Mono.just(true); + } + return null; + } + + @Override + public Mono>> getNearestMatchesAsync( + @Nonnull String collectionName, + @Nonnull Embedding embedding, + int limit, + float minRelevanceScore, + boolean withEmbedding) { + + Objects.requireNonNull(collectionName); + Objects.requireNonNull(embedding); + + if (limit <= 0) { + return Mono.just(Collections.emptyList()); + } + Query query = new Query("*=>[KNN $k @Embedding $vec AS vector_score]") + .returnFields("Id", "Text", "Description", "vector_score") + .setSortBy("vector_score", true) + .addParam("k", limit) + .addParam("vec", embeddingToBytes(embedding.getVector())) + .limit(0, limit) + .dialect(queryDialect); + String indexName = createIndexName(collectionName); + SearchResult results = this.client.ftSearch(indexName, query); + + // TODO convert to Collection> + Collection> searchTuples = + results.getDocuments().stream() + .map(document -> { + MemoryRecord memoryRecord = documentToMemoryRecord(document, withEmbedding); + return Tuples.of(memoryRecord, document.getScore().floatValue()); + }) + .collect(Collectors.toList()); + return Mono.just(searchTuples); + } + + @Override + public Mono> getNearestMatchAsync(@Nonnull String collectionName, + @Nonnull Embedding embedding, float minRelevanceScore, boolean withEmbedding) { + return getNearestMatchesAsync(collectionName, embedding, 1, minRelevanceScore, withEmbedding) + .flatMap( + nearestMatches -> { + if (nearestMatches.isEmpty()) { + return Mono.empty(); + } + return Mono.just(nearestMatches.iterator().next()); + }); + } + + /** + * Convert the Azure List embedding structure to byte[] for storage in Redis + * + * @param embeddings the list of vectors returned by the model + * @return byte[] + */ + private byte[] embeddingToBytes(List embeddings) { + ByteBuffer bytes = ByteBuffer.allocate(Float.BYTES * embeddings.size()); + bytes.order(ByteOrder.LITTLE_ENDIAN); + embeddings.iterator().forEachRemaining(bytes::putFloat); + return bytes.array(); + } + + private static OffsetDateTime ParseTimestamp(Long timestamp) { + if (timestamp != null && timestamp > 0) { + return OffsetDateTime.parse(timestamp.toString()); + } + return null; + } + + /** + * Normalize an index name + * + * @param indexName Index name + * @return A string to be used as the root of the index name; minus '-idx' + */ + private static String normalizeCollectionName(String indexName) { + if (indexName.length() > 128) { + throw new IllegalArgumentException("The collection name cannot exceed 128 chars"); + } + return indexName.toLowerCase(Locale.ROOT); + } + + private static String createIndexName(String collectionName) { + + if (collectionName.length() > 128) { + throw new IllegalArgumentException("The indexName name cannot exceed 128 chars"); + } + return String.format("%s-%s-%s", normalizeCollectionName(collectionName), RedisMemoryStore.SUFFIX, INDEX); + } + + protected Map getCollection(@Nonnull String collectionName) { + Objects.requireNonNull(collectionName); + ScanResult results = this.client.scan(collectionName); + + if (results == null) throw new MemoryException( + ErrorCodes.ATTEMPTED_TO_ACCESS_NONEXISTENT_COLLECTION, + collectionName); + + return Collections.unmodifiableMap((Map) results.getResult().stream() + .map( + key -> RedisMemoryRecord.mapToRecord( + this.client.hgetAll(String.format("%s:%s", collectionName, key)), false) + )); + } + + public static class Builder implements MemoryStore.Builder { + private String connectionString; + private int vectorSize; // in theory determined by the model but can be overridden + private VectorAlgorithm vectorIndexAlgorithm; + private String vectorDistanceMetric; + private int queryDialect = 0; // must best at least two for vector index + + public Builder connectionString(String connectionString) { + this.connectionString = connectionString; + return this; + } + public Builder vectorSize(int vectorSize) { + this.vectorSize = vectorSize; + return this; + } + public Builder vectorIndexAlgorithm(VectorAlgorithm vectorIndexAlgorithm) { + this.vectorIndexAlgorithm = vectorIndexAlgorithm; + return this; + } + public Builder vectorDistanceMetric(String vectorDistanceMetric) { + this.vectorDistanceMetric = vectorDistanceMetric; + return this; + } + + public Builder queryDialect(int queryDialect) { + this.queryDialect = queryDialect; + return this; + } + + @Override + /* + * Create a RedisMemoryStore using assigned values. + * + * @return A RedisMemoryStore object configured with the specified values + */ + public RedisMemoryStore build() { + + // make sure everything we need has been set + return new RedisMemoryStore( + connectionString, + vectorSize, + vectorIndexAlgorithm, + vectorDistanceMetric, + queryDialect + ); + } + + /** + * Create a RedisMemoryStore object using default settings; requires a valid connection string. + * + * @return A RedisMemoryStore object configured with default values + */ + public RedisMemoryStore buildDefault() { + return new RedisMemoryStore( + connectionString, + DefaultVectorSize, + RedisMemoryStore.DefaultIndexAlgorithm, + RedisMemoryStore.DefaultDistanceMetric, + RedisMemoryStore.DefaultQueryDialect + ); + } + } +} + diff --git a/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisVectorDistanceMetric.java b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisVectorDistanceMetric.java new file mode 100644 index 0000000000..f7563b793c --- /dev/null +++ b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisVectorDistanceMetric.java @@ -0,0 +1,17 @@ +package com.microsoft.semantickernel.connectors.memory.redis; + +/* + Supported distance metrics are {L2, IP, COSINE}. The default value is "COSINE". + + */ +public class RedisVectorDistanceMetric { + + /* Euclidean distance between two vectors */ + public static final String L2 = "L2"; + + /* Inner product of two vectors */ + public static final String IP = "IP"; + + /* Cosine distance of two vectors */ + public static final String COS = "COSINE"; +} diff --git a/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisVectorType.java b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisVectorType.java new file mode 100644 index 0000000000..5db1f3fbf6 --- /dev/null +++ b/java/connectors/semantickernel-connectors-memory-redis/src/main/java/com/microsoft/semantickernel/connectors/memory/redis/RedisVectorType.java @@ -0,0 +1,6 @@ +package com.microsoft.semantickernel.connectors.memory.redis; + +public class RedisVectorType { + public static final String FLOAT32 = "FLOAT32"; + public static final String FLOAT64 = "FLOAT64"; +} diff --git a/java/pom.xml b/java/pom.xml index 2aadc930d2..fe4ea2d612 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -70,6 +70,7 @@ connectors/semantickernel-connectors-memory-jdbc connectors/semantickernel-connectors-memory-postgresql connectors/semantickernel-connectors-memory-mysql + connectors/semantickernel-connectors-memory-redis extensions/semantickernel-planners extensions/semantickernel-gpt3-tokenizer