Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add document Id generator

Browse Source 
- Clean uneccessary classes/interfaces
 - straighten code style.
 - Clean tests.

 Resolves #113
This commit is contained in:
Aliakbar Jafarpour
2024-01-25 22:39:10 +01:00
committed by Christian Tzolov
parent 2704d53909
commit b0d8d2b9ce
6 changed files with 315 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
spring-ai-core/src/main/java/org/springframework/ai/document/Document.java
View File

@@ -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.
* <p>
*/
@JsonIgnoreProperties({ "contentFormatter" })
public class Document {
@@ -68,7 +74,11 @@ public class Document {
}
public Document(String content, Map<String, Object> metadata) {
this(UUID.randomUUID().toString(), content, metadata);
this(content, metadata, new RandomIdGenerator());
}
public Document(String content, Map<String, Object> metadata, IdGenerator idGenerator) {
this(idGenerator.generateId(content, metadata), content, metadata);
}
public Document(String id, String content, Map<String, Object> metadata) {

34
spring-ai-core/src/main/java/org/springframework/ai/document/id/IdGenerator.java Normal file
View File

@@ -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);
}

109
spring-ai-core/src/main/java/org/springframework/ai/document/id/JdkSha256HexIdGenerator.java Normal file
View File

@@ -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);
}
}
}

33
spring-ai-core/src/main/java/org/springframework/ai/document/id/RandomIdGenerator.java Normal file
View File

@@ -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();
}
}

67
spring-ai-core/src/test/java/org/springframework/ai/document/id/IdGeneratorProviderTest.java Normal file
View File

@@ -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<String, Object> 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<String, Object> metadata1 = Map.of("metadata", Set.of("META_DATA"));
final String content2 = content1 + " ";
final Map<String, Object> 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));
}
}

59
spring-ai-core/src/test/java/org/springframework/ai/document/id/JdkSha256HexIdGeneratorTest.java Normal file
View File

@@ -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);
}
}