Skip to content

Commit 46f12d3

Browse files
clunclun
committed
Add sample for estimated count
1 parent 2916abe commit 46f12d3

File tree

12 files changed

+370
-383
lines changed

12 files changed

+370
-383
lines changed

‎astra-db-java/src/test/java/com/datastax/astra/test/integration/collection_vectorize/AbstractVectorizeITTest.java

Lines changed: 159 additions & 156 deletions
Large diffs are not rendered by default.

‎astra-db-java/src/test/java/com/datastax/astra/test/integration/collection_vectorize/AstraDevVectorizeITTest.java

Lines changed: 54 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,7 @@
88
import com.dtsx.astra.sdk.utils.AstraEnvironment;
99
import org.junit.jupiter.api.Test;
1010

11-
import java.util.HashMap;
1211
import java.util.Map;
13-
import java.util.Optional;
14-
import java.util.function.Function;
15-
import java.util.stream.Collectors;
16-
17-
import static com.datastax.astra.test.integration.collection_vectorize.EmbeddingModel.AZURE_OPENAI_ADA002;
18-
import static com.datastax.astra.test.integration.collection_vectorize.EmbeddingModel.AZURE_OPENAI_LARGE;
19-
import static com.datastax.astra.test.integration.collection_vectorize.EmbeddingModel.AZURE_OPENAI_SMALL;
20-
import static com.datastax.astra.test.integration.collection_vectorize.EmbeddingModel.AZURE_OPENAI_SMALL_SHARED;
21-
import static org.assertj.core.api.Assertions.assertThat;
2212

2313
/**
2414
* Testing Vectorize in DEV
@@ -27,52 +17,69 @@ public class AstraDevVectorizeITTest extends AbstractVectorizeITTest {
2717

2818
@Override
2919
protected Database initDatabase() {
20+
return gcpEuropeWest4();
21+
//return gcpUsCentral();
22+
//return azureEastUs();
23+
//return awsUSWest2();
24+
}
25+
26+
private Database awsUSWest2() {
3027
Database db = initAstraDatabase(AstraEnvironment.DEV, "aws_us-west-2", CloudProviderType.AWS, "us-west-2");
3128
db.getCommandOptions().getObservers().put("logging", new LoggingCommandObserver(DataAPIClient.class));
3229
return db;
3330
}
3431

35-
@Test
36-
public void shouldTestAzureOpenAISharedSecret() {
37-
dropCollection(AZURE_OPENAI_SMALL_SHARED);
38-
String azureOpenAiKeyName = "stefano";
39-
Map<String, Object> parameters = new HashMap<>();
40-
parameters.put("deploymentId", "text-embedding-3-small-steo");
41-
parameters.put("resourceName", "steo-azure-openai");
42-
testEmbeddingModelSharedSecret(AZURE_OPENAI_SMALL_SHARED, azureOpenAiKeyName, parameters);
43-
dropCollection(AZURE_OPENAI_SMALL_SHARED);
32+
private Database gcpUsCentral() {
33+
Database db = initAstraDatabase(AstraEnvironment.DEV, "gcp_us_central1", CloudProviderType.GCP, "us-central1");
34+
db.getCommandOptions().getObservers().put("logging", new LoggingCommandObserver(DataAPIClient.class));
35+
return db;
4436
}
4537

46-
@Test
47-
public void shouldTestAzureOpenAI() {
48-
// Assessing AzureOpenAI is present
49-
Optional<EmbeddingProvider> provider = Optional.ofNullable(getDatabase()
50-
.getDatabaseAdmin()
51-
.listEmbeddingProviders()
52-
.get(AZURE_OPENAI_LARGE.getProvider()));
53-
assertThat(provider).isPresent();
54-
Map<String, EmbeddingProvider.Parameter> params = provider.get()
55-
.getParameters().stream().collect(Collectors.toMap(
56-
EmbeddingProvider.Parameter::getName,
57-
Function.identity()));
58-
assertThat(params).containsKeys("resourceName", "deploymentId");
59-
60-
String azureOpenAiKey = System.getenv("AZURE_OPENAI_API_KEY");
61-
Map<String, Object> parameters = new HashMap<>();
62-
parameters.put("deploymentId", "text-embedding-3-small-steo");
63-
parameters.put("resourceName", "steo-azure-openai");
64-
65-
dropCollection(AZURE_OPENAI_SMALL);
66-
dropCollection(AZURE_OPENAI_ADA002);
67-
dropCollection(AZURE_OPENAI_LARGE);
68-
testEmbeddingModelHeader(AZURE_OPENAI_SMALL, azureOpenAiKey, parameters);
69-
testEmbeddingModelHeader(AZURE_OPENAI_ADA002, azureOpenAiKey, parameters);
70-
testEmbeddingModelHeader(AZURE_OPENAI_LARGE, azureOpenAiKey, parameters);
71-
dropCollection(AZURE_OPENAI_SMALL);
72-
dropCollection(AZURE_OPENAI_ADA002);
73-
dropCollection(AZURE_OPENAI_LARGE);
38+
private Database gcpEuropeWest4() {
39+
Database db = initAstraDatabase(AstraEnvironment.DEV, "gcp_europe_west4", CloudProviderType.GCP, "europe-west4");
40+
db.getCommandOptions().getObservers().put("logging", new LoggingCommandObserver(DataAPIClient.class));
41+
return db;
7442
}
7543

44+
private Database azureEastUs() {
45+
Database db = initAstraDatabase(AstraEnvironment.DEV, "azure_eastus", CloudProviderType.AZURE, "eastus");
46+
db.getCommandOptions().getObservers().put("logging", new LoggingCommandObserver(DataAPIClient.class));
47+
return db;
48+
}
7649

50+
@Test
51+
public void testOneProvider() {
52+
dropAllCollections();
53+
shouldTestOneProvider("openai");
54+
shouldTestOneProvider("jinaAI");
55+
shouldTestOneProvider("voyageAI");
56+
dropAllCollections();
57+
shouldTestOneProvider("huggingface");
58+
shouldTestOneProvider("upstageAI");
59+
shouldTestOneProvider("mistral");
60+
dropAllCollections();
61+
shouldTestOneProvider("nvidia");
62+
shouldTestOneProvider("azureOpenAI");
63+
dropAllCollections();
64+
shouldTestOneProviderSharedKey("azureOpenAI", "stefano");
65+
shouldTestOneProviderSharedKey("huggingface", "HF_API_KEY");
66+
dropAllCollections();
67+
shouldTestOneProviderSharedKey("jinaAI", "JINA_API_KEY");
68+
dropAllCollections();
69+
shouldTestOneProviderSharedKey("voyageAI", "VOYAGE_API_KEY");
70+
shouldTestOneProviderSharedKey("upstageAI", "UPSTAGE_API_KEY");
71+
dropAllCollections();
72+
shouldTestOneProviderSharedKey("mistral", "MISTRAL_API_KEY");
73+
shouldTestOneProviderSharedKey("openai", "k1");
74+
}
7775

76+
@Test
77+
public void shouldTestAllProviders() {
78+
for (Map.Entry<String, EmbeddingProvider> entry : getDatabase()
79+
.getDatabaseAdmin()
80+
.listEmbeddingProviders().entrySet()) {
81+
//this.testEmbeddingProvider(entry.getKey(), entry.getValue());
82+
System.out.println("Provider: " + entry.getKey());
83+
}
84+
}
7885
}

‎astra-db-java/src/test/java/com/datastax/astra/test/integration/collection_vectorize/AstraProdVectorizeITTest.java

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
package com.datastax.astra.test.integration.collection_vectorize;
2+
3+
import com.datastax.astra.client.DataAPIClient;
4+
import com.datastax.astra.client.Database;
5+
import com.datastax.astra.client.model.EmbeddingProvider;
6+
import com.datastax.astra.internal.command.LoggingCommandObserver;
7+
import com.dtsx.astra.sdk.db.domain.CloudProviderType;
8+
import com.dtsx.astra.sdk.utils.AstraEnvironment;
9+
import org.junit.jupiter.api.Test;
10+
11+
import java.util.Map;
12+
13+
/**
14+
* Testing Vectorize in DEV
15+
*/
16+
public class AstraProdVectorizeITTest extends AbstractVectorizeITTest {
17+
18+
@Override
19+
protected Database initDatabase() {
20+
return gcpEuropeWest4();
21+
//return gcpUsEast1();
22+
//return azureEastUs();
23+
//return awsUSWest1();
24+
}
25+
26+
private Database gcpUsEast1() {
27+
Database db = initAstraDatabase(AstraEnvironment.PROD, "gcp_us_east1", CloudProviderType.GCP, "us-east1");
28+
db.getCommandOptions().getObservers().put("logging", new LoggingCommandObserver(DataAPIClient.class));
29+
return db;
30+
}
31+
private Database awsUSWest1() {
32+
Database db = initAstraDatabase(AstraEnvironment.PROD, "aws_us-west-1", CloudProviderType.AWS, "us-west-1");
33+
db.getCommandOptions().getObservers().put("logging", new LoggingCommandObserver(DataAPIClient.class));
34+
return db;
35+
}
36+
37+
38+
39+
private Database gcpEuropeWest4() {
40+
Database db = initAstraDatabase(AstraEnvironment.DEV, "gcp_europe_west4", CloudProviderType.GCP, "europe-west4");
41+
db.getCommandOptions().getObservers().put("logging", new LoggingCommandObserver(DataAPIClient.class));
42+
return db;
43+
}
44+
45+
private Database azureEastUs() {
46+
Database db = initAstraDatabase(AstraEnvironment.DEV, "azure_eastus", CloudProviderType.AZURE, "eastus");
47+
db.getCommandOptions().getObservers().put("logging", new LoggingCommandObserver(DataAPIClient.class));
48+
return db;
49+
}
50+
51+
@Test
52+
public void testOneProvider() {
53+
dropAllCollections();
54+
shouldTestOneProvider("openai");
55+
shouldTestOneProvider("jinaAI");
56+
shouldTestOneProvider("voyageAI");
57+
dropAllCollections();
58+
shouldTestOneProvider("huggingface");
59+
shouldTestOneProvider("upstageAI");
60+
shouldTestOneProvider("mistral");
61+
dropAllCollections();
62+
shouldTestOneProvider("nvidia");
63+
shouldTestOneProvider("azureOpenAI");
64+
dropAllCollections();
65+
shouldTestOneProviderSharedKey("azureOpenAI", "stefano");
66+
shouldTestOneProviderSharedKey("huggingface", "HF_API_KEY");
67+
dropAllCollections();
68+
shouldTestOneProviderSharedKey("jinaAI", "JINA_API_KEY");
69+
dropAllCollections();
70+
shouldTestOneProviderSharedKey("voyageAI", "VOYAGE_API_KEY");
71+
shouldTestOneProviderSharedKey("upstageAI", "UPSTAGE_API_KEY");
72+
dropAllCollections();
73+
shouldTestOneProviderSharedKey("mistral", "MISTRAL_API_KEY");
74+
shouldTestOneProviderSharedKey("openai", "k1");
75+
}
76+
77+
@Test
78+
public void shouldTestAllProviders() {
79+
for (Map.Entry<String, EmbeddingProvider> entry : getDatabase()
80+
.getDatabaseAdmin()
81+
.listEmbeddingProviders().entrySet()) {
82+
//this.testEmbeddingProvider(entry.getKey(), entry.getValue());
83+
System.out.println("Provider: " + entry.getKey());
84+
}
85+
}
86+
}

‎astra-db-java/src/test/java/com/datastax/astra/test/integration/collection_vectorize/EmbeddingModel.java

Lines changed: 0 additions & 86 deletions
This file was deleted.

‎astra-db-java/src/test/java/com/datastax/astra/test/integration/collection_vectorize/LocalVectorizeITTest.java

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,7 @@
66
import lombok.extern.slf4j.Slf4j;
77
import org.junit.jupiter.api.Test;
88

9-
import java.util.HashMap;
109
import java.util.Map;
11-
import java.util.Optional;
12-
import java.util.function.Function;
13-
import java.util.stream.Collectors;
14-
15-
import static com.datastax.astra.test.integration.collection_vectorize.EmbeddingModel.AZURE_OPENAI_ADA002;
16-
import static com.datastax.astra.test.integration.collection_vectorize.EmbeddingModel.AZURE_OPENAI_LARGE;
17-
import static com.datastax.astra.test.integration.collection_vectorize.EmbeddingModel.AZURE_OPENAI_SMALL;
18-
import static org.assertj.core.api.Assertions.assertThat;
1910

2011
@Slf4j
2112
public class LocalVectorizeITTest extends AbstractVectorizeITTest {
@@ -25,8 +16,23 @@ protected Database initDatabase() {
2516
return DataAPIClients.createDefaultLocalDatabase();
2617
}
2718

19+
@Test
20+
public void testOneProvider() {
21+
//shouldTestOneProvider("azureOpenAI");
22+
//shouldTestOneProvider("jinaAI");
23+
//shouldTestOneProvider("voyageAI");
24+
///shouldTestOneProvider("huggingface");
25+
//shouldTestOneProvider("upstageAI");
26+
shouldTestOneProvider("mistral");
27+
}
2828

29-
30-
29+
@Test
30+
public void shouldTestAllProviders() {
31+
for (Map.Entry<String, EmbeddingProvider> entry : getDatabase()
32+
.getDatabaseAdmin()
33+
.listEmbeddingProviders().entrySet()) {
34+
this.testEmbeddingProvider(entry.getKey(), entry.getValue());
35+
}
36+
}
3137

3238
}

0 commit comments

Comments
 (0)