Skip to content

Commit a4c7377

Browse files
clunclun
committed
Ready for release
1 parent 3cc620e commit a4c7377

File tree

90 files changed

+883
-422
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

90 files changed

+883
-422
lines changed

astra-sdk-vector/pom.xml renamed to astra-db-client/pom.xml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
<?xml version="1.0" encoding="UTF-8"?>
22
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
33
<modelVersion>4.0.0</modelVersion>
4-
<artifactId>astra-sdk-vector</artifactId>
5-
<name> + astra-sdk-vector</name>
4+
<artifactId>astra-db-client</artifactId>
5+
<name> + astra-db-client</name>
66

77
<parent>
88
<groupId>com.datastax.astra</groupId>

astra-sdk-vector/src/main/java/com/dtsx/astra/sdk/vector/VectorDatabase.java renamed to astra-db-client/src/main/java/com/dtsx/astra/sdk/AstraDB.java

Lines changed: 76 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,39 +1,78 @@
1-
package com.dtsx.astra.sdk.vector;
1+
package com.dtsx.astra.sdk;
22

33
import com.datastax.astra.sdk.AstraClient;
4-
import com.dtsx.astra.sdk.db.AstraDbClient;
4+
import com.dtsx.astra.sdk.db.AstraDBOpsClient;
55
import com.dtsx.astra.sdk.db.domain.Database;
66
import com.dtsx.astra.sdk.db.exception.DatabaseNotFoundException;
77
import com.dtsx.astra.sdk.utils.AstraEnvironment;
8-
import com.dtsx.astra.sdk.vector.domain.LLMEmbedding;
8+
import com.dtsx.astra.sdk.domain.LLMEmbedding;
9+
import io.stargate.sdk.ServiceDeployment;
10+
import io.stargate.sdk.api.SimpleTokenProvider;
11+
import io.stargate.sdk.http.ServiceHttp;
912
import io.stargate.sdk.json.JsonApiClient;
1013
import io.stargate.sdk.json.JsonCollectionClient;
1114
import io.stargate.sdk.json.JsonNamespaceClient;
1215
import io.stargate.sdk.json.domain.CollectionDefinition;
13-
import io.stargate.sdk.json.vector.JsonVectorStore;
1416
import io.stargate.sdk.json.vector.SimilarityMetric;
15-
import io.stargate.sdk.json.vector.VectorStore;
17+
import io.stargate.sdk.json.vector.VectorCollectionRepository;
18+
import io.stargate.sdk.json.vector.VectorCollectionRepositoryJson;
1619
import lombok.NonNull;
1720

21+
import java.util.Objects;
1822
import java.util.UUID;
1923
import java.util.stream.Stream;
2024

2125
/**
22-
* Client to work with a database.
26+
* Hiding top level Json Api and skip interaction with namespaces
2327
*/
24-
public class VectorDatabase {
28+
public class AstraDB {
2529

2630
/**
2731
* Hold a reference to target Astra Environment.
2832
*/
2933
protected final AstraEnvironment env;
3034

35+
/**
36+
* Top level resource for json api.
37+
*/
3138
private final JsonApiClient jsonClient;
39+
3240
/**
3341
* Namespace client
3442
*/
3543
private final JsonNamespaceClient nsClient;
3644

45+
/**
46+
* Initialization with endpoint and apikey.
47+
*
48+
* @param apiKey
49+
* api token
50+
* @param apiEndpoint
51+
* api endpoint
52+
*/
53+
public AstraDB(String apiKey, String apiEndpoint) {
54+
Objects.requireNonNull(apiKey, "apiKey");
55+
Objects.requireNonNull(apiEndpoint, "apiEndpoint");
56+
// Finding Environment based on apiEndpoint (looping to devops)
57+
if (apiEndpoint.contains(AstraEnvironment.PROD.getAppsSuffix())) {
58+
this.env = AstraEnvironment.PROD;
59+
} else if (apiEndpoint.contains(AstraEnvironment.TEST.getAppsSuffix())) {
60+
this.env = AstraEnvironment.TEST;
61+
} else if (apiEndpoint.contains(AstraEnvironment.DEV.getAppsSuffix())) {
62+
this.env = AstraEnvironment.DEV;
63+
} else {
64+
throw new IllegalArgumentException("Unable to detect environment from endpoint");
65+
}
66+
67+
// deploy on a single url with a static token (token provider)
68+
ServiceDeployment<ServiceHttp> jsonDeploy = new ServiceDeployment<>();
69+
jsonDeploy.addDatacenterTokenProvider("default", new SimpleTokenProvider(apiKey));
70+
jsonDeploy.addDatacenterServices("default", new ServiceHttp("json", apiEndpoint, apiEndpoint));
71+
this.jsonClient = new JsonApiClient(jsonDeploy);
72+
73+
this.nsClient = jsonClient.namespace("default_keyspace");
74+
}
75+
3776
/**
3877
* Full constructor.
3978
*
@@ -44,14 +83,14 @@ public class VectorDatabase {
4483
* @param env
4584
* environment
4685
*/
47-
public VectorDatabase(@NonNull String token, @NonNull UUID databaseId, @NonNull AstraEnvironment env) {
48-
this.env = env;
49-
50-
Database db = new AstraDbClient(token, env)
86+
public AstraDB(@NonNull String token, @NonNull UUID databaseId, @NonNull AstraEnvironment env) {
87+
this.env = env;
88+
Database db = new AstraDBOpsClient(token, env)
5189
.findById(databaseId.toString())
5290
.orElseThrow(() -> new DatabaseNotFoundException(databaseId.toString()));
5391

5492
this.jsonClient = AstraClient.builder()
93+
.env(env)
5594
.withDatabaseRegion(db.getInfo().getRegion())
5695
.withDatabaseId(databaseId.toString())
5796
.disableCrossRegionFailOver()
@@ -68,7 +107,7 @@ public VectorDatabase(@NonNull String token, @NonNull UUID databaseId, @NonNull
68107
* @return
69108
* name of all vector store.
70109
*/
71-
public Stream<String> findAllVectorStores() {
110+
public Stream<String> findAllCollections() {
72111
return nsClient.findCollections();
73112
}
74113

@@ -80,7 +119,7 @@ public Stream<String> findAllVectorStores() {
80119
* @return
81120
* of the store already exist
82121
*/
83-
public boolean isVectorStoreExist(@NonNull String store) {
122+
public boolean isCollectionExist(@NonNull String store) {
84123
return nsClient.existCollection(store);
85124
}
86125

@@ -90,7 +129,7 @@ public boolean isVectorStoreExist(@NonNull String store) {
90129
* @param name
91130
* store name
92131
*/
93-
public void deleteVectorStore(String name) {
132+
public void deleteCollection(String name) {
94133
nsClient.deleteCollection(name);
95134
}
96135

@@ -99,13 +138,13 @@ public void deleteVectorStore(String name) {
99138
*
100139
* @param name
101140
* store name
102-
* @param dimension
141+
* @param vectorDimension
103142
* dimension
104143
* @return
105144
* json vector store
106145
*/
107-
public JsonVectorStore createVectorStore(String name, int dimension) {
108-
nsClient.createCollection(name, dimension);
146+
public VectorCollectionRepositoryJson createCollection(String name, int vectorDimension) {
147+
nsClient.createCollection(name, vectorDimension);
109148
return nsClient.vectorStore(name);
110149
}
111150

@@ -114,13 +153,15 @@ public JsonVectorStore createVectorStore(String name, int dimension) {
114153
*
115154
* @param name
116155
* store name
117-
* @param dimension
156+
* @param vectorDimension
118157
* vector dimension
119-
* @param metric
158+
* @param vectorMetric
120159
* metric for the similarity
160+
* @return
161+
* json vector store
121162
*/
122-
public JsonVectorStore createVectorStore(String name, int dimension, SimilarityMetric metric) {
123-
nsClient.createCollection(name, dimension, metric);
163+
public VectorCollectionRepositoryJson createCollection(String name, int vectorDimension, SimilarityMetric vectorMetric) {
164+
nsClient.createCollection(name, vectorDimension, vectorMetric);
124165
return nsClient.vectorStore(name);
125166
}
126167

@@ -129,7 +170,7 @@ public JsonVectorStore createVectorStore(String name, int dimension, SimilarityM
129170
*
130171
* @param name
131172
* store name
132-
* @param dimension
173+
* @param vectorDimension
133174
* dimension
134175
* @param bean
135176
* class of pojo
@@ -138,8 +179,8 @@ public JsonVectorStore createVectorStore(String name, int dimension, SimilarityM
138179
* @param <T>
139180
* object type
140181
*/
141-
public <T> VectorStore<T> createVectorStore(String name, int dimension, Class<T> bean) {
142-
nsClient.createCollection(name, dimension);
182+
public <T> VectorCollectionRepository<T> createCollection(String name, int vectorDimension, Class<T> bean) {
183+
nsClient.createCollection(name, vectorDimension);
143184
return nsClient.vectorStore(name, bean);
144185
}
145186

@@ -148,9 +189,9 @@ public <T> VectorStore<T> createVectorStore(String name, int dimension, Class<T>
148189
*
149190
* @param name
150191
* store name
151-
* @param dimension
192+
* @param vectorDimension
152193
* dimension
153-
* @param metric
194+
* @param vectorMetric
154195
* similarity metric
155196
* @param bean
156197
* class of pojo
@@ -159,8 +200,8 @@ public <T> VectorStore<T> createVectorStore(String name, int dimension, Class<T>
159200
* @param <T>
160201
* object type
161202
*/
162-
public <T> VectorStore<T> createVectorStore(String name, int dimension, SimilarityMetric metric, Class<T> bean) {
163-
nsClient.createCollection(name, dimension, metric);
203+
public <T> VectorCollectionRepository<T> createCollection(String name, int vectorDimension, SimilarityMetric vectorMetric, Class<T> bean) {
204+
nsClient.createCollection(name, vectorDimension, vectorMetric);
164205
return nsClient.vectorStore(name, bean);
165206
}
166207

@@ -171,8 +212,10 @@ public <T> VectorStore<T> createVectorStore(String name, int dimension, Similari
171212
* vector name
172213
* @param aiModel
173214
* ai model to use
215+
* @return
216+
* json vector store
174217
*/
175-
public JsonVectorStore createVectorStore(@NonNull String name, @NonNull LLMEmbedding aiModel) {
218+
public VectorCollectionRepositoryJson createCollection(@NonNull String name, @NonNull LLMEmbedding aiModel) {
176219
nsClient.createCollection(CollectionDefinition.builder()
177220
.name(name)
178221
.vectorDimension(aiModel.getDimension())
@@ -195,7 +238,7 @@ public JsonVectorStore createVectorStore(@NonNull String name, @NonNull LLMEmbed
195238
* @return
196239
* storeName client
197240
*/
198-
public JsonVectorStore vectorStore(@NonNull String storeName) {
241+
public VectorCollectionRepositoryJson collection(@NonNull String storeName) {
199242
return nsClient.vectorStore(storeName);
200243
}
201244

@@ -211,7 +254,7 @@ public JsonVectorStore vectorStore(@NonNull String storeName) {
211254
* @param <T>
212255
* type of the bean in use
213256
*/
214-
public <T> VectorStore<T> vectorStore(@NonNull String storeName, Class<T> clazz) {
257+
public <T> VectorCollectionRepository<T> collection(@NonNull String storeName, Class<T> clazz) {
215258
return nsClient.vectorStore(storeName, clazz);
216259
}
217260

@@ -221,7 +264,7 @@ public <T> VectorStore<T> vectorStore(@NonNull String storeName, Class<T> clazz
221264
* @return
222265
* json client
223266
*/
224-
public JsonApiClient getJsonApiClient() {
267+
public JsonApiClient getRawJsonApiClient() {
225268
return jsonClient;
226269
}
227270

@@ -233,7 +276,7 @@ public JsonApiClient getJsonApiClient() {
233276
* @return
234277
* collection client
235278
*/
236-
public JsonCollectionClient rawCollectionClient(String collectionName) {
279+
public JsonCollectionClient getRawJsonCollectionClient(String collectionName) {
237280
return nsClient.collection(collectionName);
238281
}
239282

0 commit comments

Comments
 (0)