Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add tests for keyless auth to Azure Search vector store #1683

Open
wants to merge 10 commits into
base: main
Choose a base branch
from
+49 −34
Open

Add tests for keyless auth to Azure Search vector store #1683

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
c4558a9
fixing tests on Windows
mattgotteiner Nov 5, 2024
342a07d
WIP
mattgotteiner Nov 6, 2024
4e35eae
WIP
mattgotteiner Nov 6, 2024
3e3d480
WIP
mattgotteiner Nov 6, 2024
7cf7e3f
WIP
mattgotteiner Nov 6, 2024
b632d4c
WIP
mattgotteiner Nov 6, 2024
5c7f349
update
mattgotteiner Nov 6, 2024
1eeb8a1
update
mattgotteiner Nov 6, 2024
d09fbe9
update
mattgotteiner Nov 6, 2024
a5fe65b
update
mattgotteiner Nov 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,9 @@
<gemfire.testcontainers.version>2.3.0</gemfire.testcontainers.version>
<pinecone.version>0.8.0</pinecone.version>
<fastjson.version>2.0.46</fastjson.version>
<azure-core.version>1.54.0</azure-core.version>
<azure-json.version>1.3.0</azure-json.version>
<azure-identity.version>1.14.0</azure-identity.version>
<azure-search.version>11.6.1</azure-search.version>
<azure-cosmos.version>5.17.1</azure-cosmos.version>
<weaviate-client.version>4.5.1</weaviate-client.version>
Expand Down
18 changes: 18 additions & 0 deletions vector-stores/spring-ai-azure-store/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,24 @@
</exclusion>
</exclusions>
</dependency>
<!-- https://mvnrepository.com/artifact/com.azure/azure-identity -->
<dependency>
<groupId>com.azure</groupId>
<artifactId>azure-identity</artifactId>
<version>${azure-identity.version}</version>
</dependency>
<!-- https://mvnrepository.com/artifact/com.azure/azure-core -->
<dependency>
<groupId>com.azure</groupId>
<artifactId>azure-core</artifactId>
<version>${azure-core.version}</version>
</dependency>
<!-- https://mvnrepository.com/artifact/com.azure/azure-json -->
<dependency>
<groupId>com.azure</groupId>
<artifactId>azure-json</artifactId>
<version>${azure-json.version}</version>
</dependency>
<dependency>
<groupId>com.alibaba.fastjson2</groupId>
<artifactId>fastjson2</artifactId>
Expand Down
53 changes: 19 additions & 34 deletions ...org/springframework/ai/vectorstore/azure/AzureAiSearchFilterExpressionConverterTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,6 @@
*/
public class AzureAiSearchFilterExpressionConverterTests {

private static String format(String text) {
return text.trim().replace(" " + System.lineSeparator(), System.lineSeparator()) + System.lineSeparator();
}

@Test
public void testMissingFilterName() {

Expand Down Expand Up @@ -79,10 +75,9 @@ public void testEQ() {
List.of(MetadataField.text("country")));

// country == "BG"
String expected = "meta_country eq 'BG'";
String vectorExpr = converter.convertExpression(new Expression(EQ, new Key("country"), new Value("BG")));
assertThat(format(vectorExpr)).isEqualTo("""
meta_country eq 'BG'
""");
assertThat(vectorExpr).isEqualTo(expected);
}

@Test
Expand All @@ -91,12 +86,11 @@ public void tesEqAndGte() {
List.of(MetadataField.text("genre"), MetadataField.int32("year")));

// genre == "drama" AND year >= 2020
String expected = "meta_genre eq 'drama' and meta_year ge 2020";
String vectorExpr = converter
.convertExpression(new Expression(AND, new Expression(EQ, new Key("genre"), new Value("drama")),
new Expression(GTE, new Key("year"), new Value(2020))));
assertThat(format(vectorExpr)).isEqualTo("""
meta_genre eq 'drama' and meta_year ge 2020
""");
assertThat(vectorExpr).isEqualTo(expected);
}

@Test
Expand All @@ -105,11 +99,10 @@ public void tesIn() {
List.of(MetadataField.text("genre")));

// genre in ["comedy", "documentary", "drama"]
String expected = " search.in(meta_genre, 'comedy,documentary,drama', ',')";
String vectorExpr = converter.convertExpression(
new Expression(IN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(format(vectorExpr)).isEqualTo("""
search.in(meta_genre, 'comedy,documentary,drama', ',')
""");
assertThat(vectorExpr).isEqualTo(expected);
}

@Test
Expand All @@ -118,11 +111,10 @@ public void tesNin() {
List.of(MetadataField.text("genre")));

// genre in ["comedy", "documentary", "drama"]
String expected = " not search.in(meta_genre, 'comedy,documentary,drama', ',')";
String vectorExpr = converter.convertExpression(
new Expression(NIN, new Key("genre"), new Value(List.of("comedy", "documentary", "drama"))));
assertThat(format(vectorExpr)).isEqualTo("""
not search.in(meta_genre, 'comedy,documentary,drama', ',')
""");
assertThat(vectorExpr).isEqualTo(expected);
}

@Test
Expand All @@ -131,13 +123,12 @@ public void testNe() {
List.of(MetadataField.text("city"), MetadataField.int64("year"), MetadataField.text("country")));

// year >= 2020 OR country == "BG" AND city != "Sofia"
String expected = "meta_year ge 2020 or meta_country eq 'BG' and meta_city ne 'Sofia'";
String vectorExpr = converter
.convertExpression(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(AND, new Expression(EQ, new Key("country"), new Value("BG")),
new Expression(NE, new Key("city"), new Value("Sofia")))));
assertThat(format(vectorExpr)).isEqualTo("""
meta_year ge 2020 or meta_country eq 'BG' and meta_city ne 'Sofia'
""");
assertThat(vectorExpr).isEqualTo(expected);
}

@Test
Expand All @@ -146,14 +137,13 @@ public void testGroup() {
List.of(MetadataField.text("city"), MetadataField.int64("year"), MetadataField.text("country")));

// (year >= 2020 OR country == "BG") AND city != "Sofia"
String expected = "(meta_year ge 2020 or meta_country eq 'BG') and meta_city ne 'Sofia'";
String vectorExpr = converter.convertExpression(new Expression(AND,
new Group(new Expression(OR, new Expression(GTE, new Key("year"), new Value(2020)),
new Expression(EQ, new Key("country"), new Value("BG")))),
new Expression(NE, new Key("city"), new Value("Sofia"))));

assertThat(format(vectorExpr)).isEqualTo("""
(meta_year ge 2020 or meta_country eq 'BG') and meta_city ne 'Sofia'
""");
assertThat(vectorExpr).isEqualTo(expected);
}

@Test
Expand All @@ -162,14 +152,13 @@ public void tesBoolean() {
List.of(MetadataField.bool("isOpen"), MetadataField.int64("year"), MetadataField.text("country")));

// isOpen == true AND year >= 2020 AND country IN ["BG", "NL", "US"]
String expected = "meta_isOpen eq true and meta_year ge 2020 and search.in(meta_country, 'BG,NL,US', ',')";
String vectorExpr = converter.convertExpression(new Expression(AND,
new Expression(AND, new Expression(EQ, new Key("isOpen"), new Value(true)),
new Expression(GTE, new Key("year"), new Value(2020))),
new Expression(IN, new Key("country"), new Value(List.of("BG", "NL", "US")))));

assertThat(format(vectorExpr)).isEqualTo("""
meta_isOpen eq true and meta_year ge 2020 and search.in(meta_country, 'BG,NL,US', ',')
""");
assertThat(vectorExpr).isEqualTo(expected);
}

@Test
Expand All @@ -178,30 +167,26 @@ public void testDecimal() {
List.of(MetadataField.decimal("temperature")));

// temperature >= -15.6 && temperature <= +20.13
String expected = "meta_temperature ge -15.6 and meta_temperature le 20.13";
String vectorExpr = converter
.convertExpression(new Expression(AND, new Expression(GTE, new Key("temperature"), new Value(-15.6)),
new Expression(LTE, new Key("temperature"), new Value(20.13))));

assertThat(format(vectorExpr)).isEqualTo("""
meta_temperature ge -15.6 and meta_temperature le 20.13
""");
assertThat(vectorExpr).isEqualTo(expected);
}

@Test
public void testComplexIdentifiers() {
FilterExpressionConverter converter = new AzureAiSearchFilterExpressionConverter(
List.of(MetadataField.text("country 1 2 3")));

String expected = "'meta_country 1 2 3' eq 'BG'";
String vectorExpr = converter
.convertExpression(new Expression(EQ, new Key("\"country 1 2 3\""), new Value("BG")));
assertThat(format(vectorExpr)).isEqualTo("""
'meta_country 1 2 3' eq 'BG'
""");
assertThat(vectorExpr).isEqualTo(expected);

vectorExpr = converter.convertExpression(new Expression(EQ, new Key("'country 1 2 3'"), new Value("BG")));
assertThat(format(vectorExpr)).isEqualTo("""
'meta_country 1 2 3' eq 'BG'
""");
assertThat(vectorExpr).isEqualTo(expected);
}

}
9 changes: 9 additions & 0 deletions ...zure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import java.util.concurrent.TimeUnit;

import com.azure.core.credential.AzureKeyCredential;
import com.azure.identity.DefaultAzureCredentialBuilder;
import com.azure.search.documents.indexes.SearchIndexClient;
import com.azure.search.documents.indexes.SearchIndexClientBuilder;
import org.awaitility.Awaitility;
Expand Down Expand Up @@ -300,6 +301,14 @@ public static class Config {

@Bean
public SearchIndexClient searchIndexClient() {
// Only set AZURE_AI_SEARCH_TEST_KEYLESS if role-based authentication is set up correctly on the integration service
// https://learn.microsoft.com/azure/search/search-security-rbac
if (System.getenv("AZURE_AI_SEARCH_TEST_KEYLESS").equals("true")) {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe we should add some sort of useKeyless or similar name to AzureVectorStoreProperties rather than hard coding a search for an environment variable. Boot will be able to support getting this value from an env-var as well as lots of other places.

return new SearchIndexClientBuilder().endpoint(System.getenv("AZURE_AI_SEARCH_ENDPOINT"))
.credential(new DefaultAzureCredentialBuilder().build())
.buildClient();
}

return new SearchIndexClientBuilder().endpoint(System.getenv("AZURE_AI_SEARCH_ENDPOINT"))
.credential(new AzureKeyCredential(System.getenv("AZURE_AI_SEARCH_API_KEY")))
.buildClient();
Expand Down