Skip to content

Commit 1b42f5a

Browse files
clunclun
committed
demo
1 parent c2a1750 commit 1b42f5a

File tree

12 files changed

+231
-76
lines changed

12 files changed

+231
-76
lines changed

‎astra-sdk-vector/pom.xml‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,10 +35,7 @@
3535
<groupId>com.datastax.astra</groupId>
3636
<artifactId>astra-sdk-pulsar</artifactId>
3737
</exclusion>
38-
<exclusion>
39-
<groupId>com.datastax</groupId>
40-
<artifactId>astra-sdk-pulsar</artifactId>
41-
</exclusion>
38+
<!--
4239
<exclusion>
4340
<groupId>com.datastax.stargate</groupId>
4441
<artifactId>stargate-sdk-document</artifactId>
@@ -51,10 +48,10 @@
5148
<groupId>com.datastax.stargate</groupId>
5249
<artifactId>stargate-sdk-grpc</artifactId>
5350
</exclusion>
51+
-->
5452
</exclusions>
5553
</dependency>
5654

57-
5855
<!-- TEST -->
5956
<dependency>
6057
<groupId>org.junit.jupiter</groupId>

‎astra-sdk-vector/src/main/java/com/dtsx/astra/sdk/vector/AstraVectorClient.java‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public class AstraVectorClient {
5252
/**
5353
* Technical Keyspace name.
5454
*/
55-
public static final String SYSTEM_VECTOR_KEYSPACE = "vector";
55+
public static final String DEFAULT_KEYSPACE = "default_keyspace";
5656

5757
/**
5858
* First Level of the API will
@@ -146,6 +146,22 @@ public UUID createDatabase(@NonNull String name) {
146146
return createDatabase(name, FREE_TIER_CLOUD, FREE_TIER_CLOUD_REGION);
147147
}
148148

149+
/**
150+
* Create Db in free Tier.
151+
*
152+
* @param name
153+
* database name
154+
* @return
155+
* if the db has been deleted
156+
*/
157+
public boolean deleteDatabase(@NonNull String name) {
158+
Optional<Database> opDb = findByName(name).findFirst();
159+
opDb.ifPresent(db -> {
160+
devopsDbClient.database(db.getId()).delete();
161+
});
162+
return opDb.isPresent();
163+
}
164+
149165
/**
150166
* Create a database with the full definition.
151167
*
@@ -188,7 +204,7 @@ public UUID createDatabase(@NonNull String name, @NonNull CloudProviderType clou
188204
.name(name)
189205
.cloudProvider(cloud)
190206
.cloudRegion(cloudRegion)
191-
.keyspace(SYSTEM_VECTOR_KEYSPACE)
207+
.keyspace(DEFAULT_KEYSPACE)
192208
.withVector().build()));
193209
log.info("Database {} is starting (id={}): it will take about a minute please wait...", name, newDbId);
194210
waitForDatabase(devopsDbClient.database(newDbId.toString()));
@@ -242,7 +258,7 @@ public Optional<Database> findById(@NonNull UUID id) {
242258
* @return
243259
* database client
244260
*/
245-
public AstraVectorDatabaseClient database(@NonNull String databaseName) {
261+
public VectorDatabase database(@NonNull String databaseName) {
246262
List<Database> dbs = findByName(databaseName).collect(Collectors.toList());
247263
if (dbs.isEmpty()) {
248264
throw new DatabaseNotFoundException(databaseName);
@@ -265,8 +281,8 @@ public AstraVectorDatabaseClient database(@NonNull String databaseName) {
265281
* @return
266282
* database client
267283
*/
268-
public AstraVectorDatabaseClient database(UUID databaseId) {
269-
return new AstraVectorDatabaseClient(token, databaseId, env);
284+
public VectorDatabase database(UUID databaseId) {
285+
return new VectorDatabase(token, databaseId, env);
270286
}
271287

272288
/**

astra-sdk-vector/src/main/java/com/dtsx/astra/sdk/vector/AstraVectorDatabaseClient.java renamed to astra-sdk-vector/src/main/java/com/dtsx/astra/sdk/vector/VectorDatabase.java

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
import com.dtsx.astra.sdk.db.exception.DatabaseNotFoundException;
77
import com.dtsx.astra.sdk.utils.AstraEnvironment;
88
import com.dtsx.astra.sdk.vector.domain.LLMEmbedding;
9-
import io.stargate.sdk.core.domain.ObjectMap;
109
import io.stargate.sdk.json.JsonApiClient;
1110
import io.stargate.sdk.json.JsonCollectionClient;
1211
import io.stargate.sdk.json.JsonNamespaceClient;
@@ -22,13 +21,14 @@
2221
/**
2322
* Client to work with a database.
2423
*/
25-
public class AstraVectorDatabaseClient {
24+
public class VectorDatabase {
2625

2726
/**
2827
* Hold a reference to target Astra Environment.
2928
*/
3029
protected final AstraEnvironment env;
3130

31+
private final JsonApiClient jsonClient;
3232
/**
3333
* Namespace client
3434
*/
@@ -44,20 +44,21 @@ public class AstraVectorDatabaseClient {
4444
* @param env
4545
* environment
4646
*/
47-
public AstraVectorDatabaseClient(@NonNull String token, @NonNull UUID databaseId, @NonNull AstraEnvironment env) {
47+
public VectorDatabase(@NonNull String token, @NonNull UUID databaseId, @NonNull AstraEnvironment env) {
4848
this.env = env;
4949

5050
Database db = new AstraDbClient(token, env)
5151
.findById(databaseId.toString())
5252
.orElseThrow(() -> new DatabaseNotFoundException(databaseId.toString()));
5353

54-
JsonApiClient jsonClient = AstraClient.builder()
54+
this.jsonClient = AstraClient.builder()
5555
.withDatabaseRegion(db.getInfo().getRegion())
5656
.withDatabaseId(databaseId.toString())
5757
.disableCrossRegionFailOver()
5858
.build()
5959
.apiStargateJson();
6060

61+
// will inherit 'default_keyspace' from the database
6162
this.nsClient = jsonClient.namespace(db.getInfo().getKeyspace());
6263
}
6364

@@ -154,7 +155,7 @@ public void createVectorStore(@NonNull String name, @NonNull LLMEmbedding aiMode
154155
.name(name)
155156
.vectorDimension(aiModel.getDimension())
156157
.similarityMetric(SimilarityMetric.cosine)
157-
.llmProvider(aiModel.getLlmprovider().name())
158+
.llmProvider(aiModel.getLlmprovider())
158159
.llmModel(aiModel.getName())
159160
.build());
160161
}
@@ -191,6 +192,16 @@ public <T> VectorStore<T> vectorStore(@NonNull String storeName, Class<T> clazz
191192
return nsClient.vectorStore(storeName, clazz);
192193
}
193194

195+
/**
196+
* Access json api client
197+
*
198+
* @return
199+
* json client
200+
*/
201+
public JsonApiClient getJsonApiClient() {
202+
return jsonClient;
203+
}
204+
194205
/**
195206
* Access to low level object.
196207
*

‎astra-sdk-vector/src/main/java/com/dtsx/astra/sdk/vector/domain/LLMEmbedding.java‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -7,24 +7,18 @@
77
public enum LLMEmbedding {
88

99
/** openai. */
10-
OPENAI_TEXT_DAVINCI_003(LLmProvider.openAI,"text-davinci-003", 1536),
11-
12-
/** openai. */
13-
OPENAI_TEXT_EMBEDDING_ADA_002(LLmProvider.openAI,"text-embedding-ada-002", 1536),
10+
ADA_002("openai","text-embedding-ada-002", 1536),
1411

1512
/** google. */
16-
PALM2(LLmProvider.google,"text-bison", 768),
13+
GECKO("google", "textembedding-gecko", 768),
1714

18-
/** google. */
19-
CODEY(LLmProvider.google,"chat-bison", 768),
20-
21-
/** google. */
22-
VERTEX_AI(LLmProvider.google, "textembedding-gecko", 768);
15+
/** hf. */
16+
MINILM_L6_V2("hugging-face", "all-MiniLM-L6-v2", 384);
2317

2418
/**
2519
* llm provider
2620
*/
27-
final LLmProvider llmprovider;
21+
final String llmprovider;
2822

2923
/**
3024
* model name
@@ -45,7 +39,7 @@ public enum LLMEmbedding {
4539
* @param dimension
4640
* model dimension
4741
*/
48-
LLMEmbedding(LLmProvider provider, String name, int dimension) {
42+
LLMEmbedding(String provider, String name, int dimension) {
4943
this.llmprovider = provider;
5044
this.name = name;
5145
this.dimension = dimension;

‎astra-sdk-vector/src/main/java/com/dtsx/astra/sdk/vector/domain/LLmProvider.java‎

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

‎astra-sdk-vector/src/test/java/com/dtsx/astra/sdk/vector/AstraVectorQuickStart.java‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
package com.dtsx.astra.sdk.vector;
22

3+
import com.dtsx.astra.sdk.db.domain.Database;
4+
import com.dtsx.astra.sdk.db.domain.DatabaseInfo;
35
import com.fasterxml.jackson.annotation.JsonProperty;
6+
import io.stargate.sdk.core.domain.ObjectMap;
47
import io.stargate.sdk.json.domain.JsonDocument;
58
import io.stargate.sdk.json.domain.odm.Document;
9+
import io.stargate.sdk.json.domain.odm.Result;
610
import io.stargate.sdk.json.vector.JsonVectorStore;
711
import io.stargate.sdk.json.vector.VectorStore;
812
import lombok.AllArgsConstructor;
@@ -11,13 +15,11 @@
1115
import org.junit.jupiter.api.Assertions;
1216
import org.junit.jupiter.api.Test;
1317

18+
import java.util.List;
1419
import java.util.Map;
1520

1621
public class AstraVectorQuickStart {
1722

18-
19-
20-
2123
@Test
2224
public void quickStartTest() {
2325
String databaseName = "vector_client_test";
@@ -27,13 +29,18 @@ public void quickStartTest() {
2729
// 1a. Initialization with a client
2830
AstraVectorClient vectorClient = new AstraVectorClient(astraToken);
2931

32+
vectorClient
33+
.findAllDatabases()
34+
.map(Database::getInfo)
35+
.map(DatabaseInfo::getName)
36+
.forEach(System.out::println);
3037
// 1b. Create DB (Skip if you already have a database running)
3138
if (!vectorClient.isDatabaseExists(databaseName)) {
3239
vectorClient.createDatabase(databaseName);
3340
}
3441

3542
// 2. Create a store (delete if exist)
36-
AstraVectorDatabaseClient vectorDb = vectorClient.database(databaseName);
43+
VectorDatabase vectorDb = vectorClient.database(databaseName);
3744
vectorDb.deleteStore(vectorStoreName);
3845
vectorDb.createVectorStore(vectorStoreName, 14);
3946

@@ -83,6 +90,10 @@ public void quickStartTest() {
8390
productVectorStore.insert(doc6);
8491

8592
Assertions.assertEquals(6, productVectorStore.count());
93+
94+
List<JsonDocument> results = vectorStore
95+
.similaritySearchJson(new float[]{1f, 1f, 1f, 1f, 1f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f}, 2);
96+
8697
}
8798

8899

‎astra-sdk-vector/src/test/java/com/dtsx/astra/sdk/vector/VectorClientDbCreationsTest.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,8 @@ public void shouldCreateDatabases() {
8080
AstraVectorClient.FREE_TIER_CLOUD,
8181
AstraVectorClient.FREE_TIER_CLOUD_REGION);
8282
log.info("{} is created and active", yaDatabaseId);
83+
84+
vectorClient.devopsDbClient.database("").delete();
8385
}
8486

8587
@Test
@@ -91,14 +93,14 @@ public void shouldCreateCollection() {
9193
AstraVectorClient vectorClient = new AstraVectorClient();
9294

9395
// Assign vectorDb client (but could user fluent everywhere)
94-
AstraVectorDatabaseClient vectorDb = vectorClient.database(DBNAME_VECTOR_CLIENT);
96+
VectorDatabase vectorDb = vectorClient.database(DBNAME_VECTOR_CLIENT);
9597

9698
// Create a Store with $vector
9799
vectorDb.createVectorStore(DEMO_COLLECTION_VECTOR, 14);
98100
log.info("{} is created", DEMO_COLLECTION_VECTOR);
99101

100102
// [NOT FOR GA] -Create a Store with $vectorize
101-
vectorDb.createVectorStore(STORE_WITH_VECTORIZE_TMP, LLMEmbedding.OPENAI_TEXT_EMBEDDING_ADA_002);
103+
vectorDb.createVectorStore(STORE_WITH_VECTORIZE_TMP, LLMEmbedding.ADA_002);
102104
log.info("{} is created", STORE_WITH_VECTORIZE_TMP);
103105

104106
log.info("List of Stores");
@@ -111,7 +113,7 @@ public void shouldCreateCollection() {
111113
@EnabledIfEnvironmentVariable(named = "ASTRA_DB_APPLICATION_TOKEN", matches = "Astra.*")
112114
public void shouldDeleteStore() {
113115
// Assign vectorDb client (but could user fluent everywhere)
114-
AstraVectorDatabaseClient vectorDb = new AstraVectorClient().database(DBNAME_VECTOR_CLIENT);
116+
VectorDatabase vectorDb = new AstraVectorClient().database(DBNAME_VECTOR_CLIENT);
115117
// Given
116118
Assertions.assertTrue(vectorDb.isStoreExist(STORE_WITH_VECTORIZE_TMP));
117119
// When Delete a store

‎astra-sdk-vector/src/test/java/com/dtsx/astra/sdk/vector/VectorClientPhilosopherTest.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,7 @@
33
import com.dtsx.astra.sdk.utils.AstraRc;
44
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
55
import dev.langchain4j.model.openai.OpenAiModelName;
6-
import io.stargate.sdk.json.JsonApiClient;
76
import io.stargate.sdk.json.domain.odm.Document;
8-
import io.stargate.sdk.json.vector.SimilarityMetric;
97
import io.stargate.sdk.json.vector.VectorStore;
108
import lombok.AllArgsConstructor;
119
import lombok.Data;
@@ -75,7 +73,7 @@ public static void setup() {
7573
@EnabledIfEnvironmentVariable(named = "ASTRA_DB_APPLICATION_TOKEN", matches = "Astra.*")
7674
public void shouldIngestCsv() {
7775
// Init the Store
78-
AstraVectorDatabaseClient dbClient = new AstraVectorClient().database(DBNAME_VECTOR_CLIENT);
76+
VectorDatabase dbClient = new AstraVectorClient().database(DBNAME_VECTOR_CLIENT);
7977
dbClient.deleteStore(VECTOR_STORE_NAME);
8078
vectorStore = dbClient.createVectorStore(VECTOR_STORE_NAME, 1536, Quote.class);
8179
log.info("store {} is created ", VECTOR_STORE_NAME);

‎astra-sdk-vector/src/test/java/com/dtsx/astra/sdk/vector/VectorClientProductTest.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import com.dtsx.astra.sdk.utils.AstraRc;
44
import com.fasterxml.jackson.annotation.JsonProperty;
5-
import io.stargate.sdk.core.domain.Page;
65
import io.stargate.sdk.json.domain.Filter;
76
import io.stargate.sdk.json.domain.JsonDocument;
87
import io.stargate.sdk.json.domain.odm.Document;
@@ -12,7 +11,6 @@
1211
import lombok.Data;
1312
import lombok.NoArgsConstructor;
1413
import lombok.extern.slf4j.Slf4j;
15-
import org.junit.jupiter.api.Assertions;
1614
import org.junit.jupiter.api.BeforeAll;
1715
import org.junit.jupiter.api.DisplayName;
1816
import org.junit.jupiter.api.Order;
@@ -65,7 +63,7 @@ public static void setup() {
6563
@EnabledIfEnvironmentVariable(named = "ASTRA_DB_APPLICATION_TOKEN", matches = "Astra.*")
6664
public void shouldInsertStaticDocument() {
6765
// Recreating the store
68-
AstraVectorDatabaseClient dbClient = new AstraVectorClient().database(DBNAME_VECTOR_CLIENT);
66+
VectorDatabase dbClient = new AstraVectorClient().database(DBNAME_VECTOR_CLIENT);
6967
dbClient.deleteStore(VECTOR_STORE_NAME);
7068
vectorStore = dbClient.createVectorStore(VECTOR_STORE_NAME, 14, Product.class);
7169
log.info("store {} is created ", VECTOR_STORE_NAME);

0 commit comments

Comments
 (0)