diff --git a/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java b/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java index 8554a1081..078750af1 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java @@ -1,5 +1,5 @@ /* - * Copyright 2023-2023 the original author or authors. + * Copyright 2023-2024 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -20,15 +20,21 @@ import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.UUID; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import com.fasterxml.jackson.annotation.JsonProperty; +import org.springframework.ai.document.id.IdGenerator; +import org.springframework.ai.document.id.RandomIdGenerator; import org.springframework.util.Assert; +/** + * A document is a container for the content and metadata of a document. It also contains + * the document's unique ID and an optional embedding. + *

+ */ @JsonIgnoreProperties({ "contentFormatter" }) public class Document { @@ -68,7 +74,11 @@ public class Document { } public Document(String content, Map metadata) { - this(UUID.randomUUID().toString(), content, metadata); + this(content, metadata, new RandomIdGenerator()); + } + + public Document(String content, Map metadata, IdGenerator idGenerator) { + this(idGenerator.generateId(content, metadata), content, metadata); } public Document(String id, String content, Map metadata) { diff --git a/spring-ai-core/src/main/java/org/springframework/ai/document/id/IdGenerator.java b/spring-ai-core/src/main/java/org/springframework/ai/document/id/IdGenerator.java new file mode 100644 index 000000000..1aab8841c --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/document/id/IdGenerator.java @@ -0,0 +1,34 @@ +/* + * Copyright 2024-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.document.id; + +/** + * Interface for generating unique document IDs. + * + * @author Aliakbar Jafarpour + * @author Christian Tzolov + */ +public interface IdGenerator { + + /** + * Generate a unique ID for the given content. Note: some generator, such as the the + * random generator might not dependant on or use the content parameters. + * @param contents the content to generate an ID for. + * @return the generated ID. + */ + String generateId(Object... contents); + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/document/id/JdkSha256HexIdGenerator.java b/spring-ai-core/src/main/java/org/springframework/ai/document/id/JdkSha256HexIdGenerator.java new file mode 100644 index 000000000..08f34aea5 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/document/id/JdkSha256HexIdGenerator.java @@ -0,0 +1,109 @@ +/* + * Copyright 2024-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.document.id; + +import java.io.ByteArrayOutputStream; +import java.io.ObjectOutputStream; +import java.nio.charset.Charset; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.UUID; + +import org.springframework.util.Assert; + +/** + * A SHA-256 based ID generator that returns the hash as a UUID. + * + * @author Aliakbar Jafarpour + * @author Christian Tzolov + */ +public class JdkSha256HexIdGenerator implements IdGenerator { + + private static final String SHA_256 = "SHA-256"; + + private final String byteHexFormat = "%02x"; + + private final Charset charset; + + private final MessageDigest messageDigest; + + public JdkSha256HexIdGenerator(final String algorithm, final Charset charset) { + this.charset = charset; + try { + this.messageDigest = MessageDigest.getInstance(algorithm); + } + catch (NoSuchAlgorithmException e) { + throw new IllegalArgumentException(e); + } + } + + public JdkSha256HexIdGenerator() { + this(SHA_256, StandardCharsets.UTF_8); + } + + @Override + public String generateId(Object... contents) { + return this.hash(this.serializeToBytes(contents)); + } + + // https://github.com/spring-projects/spring-ai/issues/113#issue-2000373318 + private String hash(byte[] contentWithMetadata) { + byte[] hashBytes = getMessageDigest().digest(contentWithMetadata); + StringBuilder sb = new StringBuilder(); + for (byte b : hashBytes) { + sb.append(String.format(this.byteHexFormat, b)); + } + return UUID.nameUUIDFromBytes(sb.toString().getBytes(this.charset)).toString(); + } + + private byte[] serializeToBytes(Object... contents) { + Assert.notNull(contents, "Contents must not be null"); + ByteArrayOutputStream byteOut = null; + try { + byteOut = new ByteArrayOutputStream(); + ObjectOutputStream out = new ObjectOutputStream(byteOut); + for (Object content : contents) { + out.writeObject(content); + } + return byteOut.toByteArray(); + } + catch (Exception e) { + throw new RuntimeException("Failed to serialize", e); + } + finally { + if (byteOut != null) { + try { + byteOut.close(); + } + catch (Exception e) { + // ignore + } + } + } + } + + MessageDigest getMessageDigest() { + try { + return (MessageDigest) messageDigest.clone(); + } + catch (CloneNotSupportedException e) { + throw new RuntimeException("Unsupported clone for MessageDigest.", e); + } + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/document/id/RandomIdGenerator.java b/spring-ai-core/src/main/java/org/springframework/ai/document/id/RandomIdGenerator.java new file mode 100644 index 000000000..c7006357c --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/document/id/RandomIdGenerator.java @@ -0,0 +1,33 @@ +/* + * Copyright 2024-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.document.id; + +import java.util.UUID; + +/** + * A random ID generator that returns a UUID. + * + * @author Aliakbar Jafarpour + * @author Christian Tzolov + */ +public class RandomIdGenerator implements IdGenerator { + + @Override + public String generateId(Object... contents) { + return UUID.randomUUID().toString(); + } + +} diff --git a/spring-ai-core/src/test/java/org/springframework/ai/document/id/IdGeneratorProviderTest.java b/spring-ai-core/src/test/java/org/springframework/ai/document/id/IdGeneratorProviderTest.java new file mode 100644 index 000000000..874311264 --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/document/id/IdGeneratorProviderTest.java @@ -0,0 +1,67 @@ +/* + * Copyright 2024-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.document.id; + +import java.util.Map; +import java.util.Set; +import java.util.UUID; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class IdGeneratorProviderTest { + + @Test + void hashGeneratorGenerateSimilarIdsForSimilarContent() { + + var idGenerator1 = new JdkSha256HexIdGenerator(); + var idGenerator2 = new JdkSha256HexIdGenerator(); + + final String content = "Content"; + final Map metadata = Map.of("metadata", Set.of("META_DATA")); + + String actualHashes1 = idGenerator1.generateId(content, metadata); + String actualHashes2 = idGenerator2.generateId(content, metadata); + + Assertions.assertEquals(actualHashes1, actualHashes2); + + // Assert (other expected behaviors) + Assertions.assertDoesNotThrow(() -> UUID.fromString(actualHashes1)); + Assertions.assertDoesNotThrow(() -> UUID.fromString(actualHashes2)); + } + + @Test + void hashGeneratorGenerateDifferentIdsForDifferentContent() { + + var idGenerator1 = new JdkSha256HexIdGenerator(); + var idGenerator2 = new JdkSha256HexIdGenerator(); + + final String content1 = "Content"; + final Map metadata1 = Map.of("metadata", Set.of("META_DATA")); + final String content2 = content1 + " "; + final Map metadata2 = metadata1; + + String actualHashes1 = idGenerator1.generateId(content1, metadata1); + String actualHashes2 = idGenerator2.generateId(content2, metadata2); + + Assertions.assertNotEquals(actualHashes1, actualHashes2); + + // Assert (other expected behaviors) + Assertions.assertDoesNotThrow(() -> UUID.fromString(actualHashes1)); + Assertions.assertDoesNotThrow(() -> UUID.fromString(actualHashes2)); + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/test/java/org/springframework/ai/document/id/JdkSha256HexIdGeneratorTest.java b/spring-ai-core/src/test/java/org/springframework/ai/document/id/JdkSha256HexIdGeneratorTest.java new file mode 100644 index 000000000..95ab69482 --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/document/id/JdkSha256HexIdGeneratorTest.java @@ -0,0 +1,59 @@ +/* + * Copyright 2024-2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.document.id; + +import java.nio.charset.Charset; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; + +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.Test; + +public class JdkSha256HexIdGeneratorTest { + + private final JdkSha256HexIdGenerator testee = new JdkSha256HexIdGenerator(); + + @Test + void messageDigestReturnsDistinctInstances() { + final MessageDigest md1 = testee.getMessageDigest(); + final MessageDigest md2 = testee.getMessageDigest(); + + Assertions.assertThat(md1 != md2).isTrue(); + + Assertions.assertThat(md1.getAlgorithm()).isEqualTo(md2.getAlgorithm()); + Assertions.assertThat(md1.getDigestLength()).isEqualTo(md2.getDigestLength()); + Assertions.assertThat(md1.getProvider()).isEqualTo(md2.getProvider()); + Assertions.assertThat(md1.toString()).isEqualTo(md2.toString()); + } + + @Test + void messageDigestReturnsInstancesWithIndependentAndReproducibleDigests() { + final String updateString1 = "md1_update"; + final String updateString2 = "md2_update"; + final Charset charset = StandardCharsets.UTF_8; + + final byte[] md1BytesFirstTry = testee.getMessageDigest().digest(updateString1.getBytes(charset)); + final byte[] md2BytesFirstTry = testee.getMessageDigest().digest(updateString2.getBytes(charset)); + final byte[] md1BytesSecondTry = testee.getMessageDigest().digest(updateString1.getBytes(charset)); + final byte[] md2BytesSecondTry = testee.getMessageDigest().digest(updateString2.getBytes(charset)); + + Assertions.assertThat(md1BytesFirstTry).isNotEqualTo(md2BytesFirstTry); + + Assertions.assertThat(md1BytesFirstTry).isEqualTo(md1BytesSecondTry); + Assertions.assertThat(md2BytesFirstTry).isEqualTo(md2BytesSecondTry); + } + +} \ No newline at end of file