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

Harden milvus IT test fixtures

Browse Source
This commit is contained in:
Christian Tzolov
2023-09-28 05:25:22 +02:00
parent 5ff2c3bb92
commit a45e34edd7
2 changed files with 27 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/MilvusVectorStore.java
View File

@@ -48,8 +48,8 @@ import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.param.index.DescribeIndexParam;
import io.milvus.param.index.DropIndexParam;
import io.milvus.response.SearchResultsWrapper;
import io.milvus.response.QueryResultsWrapper.RowRecord;
import io.milvus.response.SearchResultsWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -359,7 +359,7 @@ public class MilvusVectorStore implements VectorStore, SmartLifecycle {
metadata.put(DISTANCE_FIELD_NAME, 1 - getResultSimilarity(rowRecord));
return new Document(docId, content, metadata.getInnerMap());
})
.collect(Collectors.toList());
.toList();
}
private float getResultSimilarity(RowRecord rowRecord) {
@@ -368,7 +368,7 @@ public class MilvusVectorStore implements VectorStore, SmartLifecycle {
}
private List<Float> toFloatList(List<Double> embeddingDouble) {
return embeddingDouble.stream().map(Number::floatValue).collect(Collectors.toList());
return embeddingDouble.stream().map(Number::floatValue).toList();
}
// ---------------------------------------------------------------------------------

38
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/MilvusVectorStoreIT.java
View File

@@ -21,18 +21,17 @@ import java.time.Duration;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
import org.testcontainers.containers.DockerComposeContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration;
@@ -56,14 +55,9 @@ import static org.assertj.core.api.Assertions.assertThat;
@Testcontainers
public class MilvusVectorStoreIT {
@Container
public static DockerComposeContainer milvusContainer = new DockerComposeContainer(
new File("src/test/resources/docker-compose.yml"))
.withExposedService("standalone", 19530)
.withExposedService("standalone", 9091,
Wait.forHttp("/healthz").forPort(9091).forStatusCode(200).forStatusCode(401))
.waitingFor("standalone",
Wait.forLogMessage(".*Proxy successfully started.*\\s", 1).withStartupTimeout(Duration.ofSeconds(100)));
private static DockerComposeContainer milvusContainer;
private static final File TEMP_FOLDER = new File("target/test-" + UUID.randomUUID().toString());
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withUserConfiguration(TestApplication.class)
@@ -77,9 +71,25 @@ public class MilvusVectorStoreIT {
"Great Depression Great Depression Great Depression Great Depression Great Depression Great Depression",
Collections.singletonMap("meta2", "meta2")));
@BeforeAll
public static void beforeAll() {
FileSystemUtils.deleteRecursively(TEMP_FOLDER);
TEMP_FOLDER.mkdirs();
milvusContainer = new DockerComposeContainer(new File("src/test/resources/docker-compose.yml"))
.withEnv("DOCKER_VOLUME_DIRECTORY", TEMP_FOLDER.getAbsolutePath())
.withExposedService("standalone", 19530)
.withExposedService("standalone", 9091,
Wait.forHttp("/healthz").forPort(9091).forStatusCode(200).forStatusCode(401))
.waitingFor("standalone", Wait.forLogMessage(".*Proxy successfully started.*\\s", 1)
.withStartupTimeout(Duration.ofSeconds(100)));
milvusContainer.start();
}
@AfterAll
public static void afterAll() {
FileSystemUtils.deleteRecursively(new File("src/test/resources/volumes"));
milvusContainer.stop();
FileSystemUtils.deleteRecursively(TEMP_FOLDER);
}
private void resetCollection(VectorStore vectorStore) {
@@ -113,7 +123,7 @@ public class MilvusVectorStoreIT {
assertThat(resultDoc.getMetadata()).containsKey("distance");
// Remove all documents from the store
vectorStore.delete(documents.stream().map(doc -> doc.getId()).collect(Collectors.toList()));
vectorStore.delete(documents.stream().map(doc -> doc.getId()).toList());
List<Document> results2 = vectorStore.similaritySearch("Hello", 1);
assertThat(results2).hasSize(0);
@@ -184,11 +194,11 @@ public class MilvusVectorStoreIT {
List<Float> distances = fullResult.stream()
.map(doc -> (Float) doc.getMetadata().get("distance"))
.collect(Collectors.toList());
.toList();
assertThat(distances).hasSize(3);
List<Document> results = vectorStore.similaritySearch("Great", 5, (1 - (distances.get(0) + 0.01)));
List<Document> results = vectorStore.similaritySearch("Great", 5, (1 - (distances.get(0) + 0.001)));
assertThat(results).hasSize(1);
Document resultDoc = results.get(0);