Skip to content

Commit

Permalink
fixing IT
Browse files Browse the repository at this point in the history
  • Loading branch information
langchain4j committed May 24, 2024
1 parent a661050 commit 5b4b9f1
Showing 1 changed file with 33 additions and 10 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import com.azure.search.documents.indexes.models.SearchIndex;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.internal.Utils;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.Content;
Expand All @@ -26,7 +27,6 @@

import java.util.List;

import static dev.langchain4j.store.embedding.azure.search.AbstractAzureAiSearchEmbeddingStore.DEFAULT_INDEX_NAME;
import static java.util.Arrays.asList;
import static org.assertj.core.api.Assertions.assertThat;

Expand All @@ -44,26 +44,21 @@ class AutoConfigIT {
.endpoint(System.getenv("AZURE_SEARCH_ENDPOINT"))
.credential(new AzureKeyCredential(System.getenv("AZURE_SEARCH_KEY")))
.buildClient();
private final SearchIndex index = new SearchIndex(DEFAULT_INDEX_NAME);

private final String indexName = Utils.randomUUID();
private final SearchIndex index = new SearchIndex(indexName);

private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(AutoConfig.class));

@AfterEach
void afterEach() {
searchIndexClient.deleteIndex(DEFAULT_INDEX_NAME);
try {
Thread.sleep(10_000);
} catch (InterruptedException e) {
throw new RuntimeException(e);
}
searchIndexClient.deleteIndex(indexName);
}

@Test
void should_provide_ai_search_retriever() {

// TODO extract in separate tests

contextRunner
.withPropertyValues(
Properties.PREFIX + ".content-retriever.api-key=" + AZURE_SEARCH_KEY,
Expand Down Expand Up @@ -159,6 +154,34 @@ void should_provide_ai_search_retriever() {
@EnabledIfEnvironmentVariable(named = "AZURE_SEARCH_RERANKER_AVAILABLE", matches = "true")
void should_provide_ai_search_retriever_with_reranking() {

contextRunner
.withPropertyValues(
Properties.PREFIX + ".content-retriever.api-key=" + AZURE_SEARCH_KEY,
Properties.PREFIX + ".content-retriever.endpoint=" + AZURE_SEARCH_ENDPOINT,
Properties.PREFIX + ".content-retriever.dimensions=" + dimensions,
Properties.PREFIX + ".content-retriever.create-or-update-index=" + "true",
Properties.PREFIX + ".content-retriever.query-type=" + "VECTOR"
).withBean(EmbeddingModel.class, () -> embeddingModel)
.run(context -> {
ContentRetriever contentRetriever = context.getBean(ContentRetriever.class);
assertThat(contentRetriever).isInstanceOf(AzureAiSearchContentRetriever.class);
AzureAiSearchContentRetriever azureAiSearchContentRetriever = (AzureAiSearchContentRetriever) contentRetriever;

String content1 = "This book is about politics";
String content2 = "Cats sleeps a lot.";
String content3 = "Sandwiches taste good.";
String content4 = "The house is open";
List<String> contents = asList(content1, content2, content3, content4);

for (String content : contents) {
TextSegment textSegment = TextSegment.from(content);
Embedding embedding = embeddingModel.embed(content).content();
azureAiSearchContentRetriever.add(embedding, textSegment);
}

awaitUntilPersisted();
});

String content = "house";
Query query = Query.from(content);

Expand Down

0 comments on commit 5b4b9f1

Please sign in to comment.