Skip to content

Commit f17e6f1

Browse files
clunclun
committed
Astra Vector Client
1 parent 0864390 commit f17e6f1

File tree

3 files changed

+71
-32
lines changed

3 files changed

+71
-32
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,7 @@ public void createVectorStore(@NonNull String name, @NonNull LLMEmbedding aiMode
153153
* @return
154154
* storeName client
155155
*/
156-
public VectorStore<ObjectMap> vectorStore(@NonNull String storeName) {
156+
public JsonVectorStore vectorStore(@NonNull String storeName) {
157157
return nsClient.vectorStore(storeName);
158158
}
159159

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

Lines changed: 69 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -10,54 +10,92 @@
1010
import lombok.AllArgsConstructor;
1111
import lombok.Data;
1212
import lombok.NoArgsConstructor;
13+
import org.junit.jupiter.api.Assertions;
14+
import org.junit.jupiter.api.Test;
1315

1416
import java.util.Map;
1517
import java.util.UUID;
1618

1719
public class AstraVectorQuickStart {
1820

21+
@Test
22+
public void theCountLoop() {
23+
String databaseName = "vector_client_test";
24+
String vectorStoreName = "collection_count";
25+
String astraToken = System.getenv("ASTRA_DB_APPLICATION_TOKEN");
26+
27+
AstraVectorClient vectorClient = new AstraVectorClient(astraToken);
28+
AstraVectorDatabaseClient vectorDb = vectorClient.database(databaseName);
29+
//vectorDb.createVectorStore(vectorStoreName, 14);
30+
JsonVectorStore vectorStore = vectorDb.vectorStore(vectorStoreName);
31+
32+
for(int i=vectorStore.count();i<1000;i++) {
33+
vectorStore.insert(new JsonDocument()
34+
.vector(new float[]{(float) Math.random(), (float) Math.random(), (float) Math.random(), (float) Math.random(), (float) Math.random(), (float) Math.random(),(float) Math.random(), (float) Math.random(),(float) Math.random(), (float) Math.random(), (float) Math.random(), (float) Math.random(), (float) Math.random(), (float) Math.random()})
35+
.put("product_name", "HealthyFresh - Beef raw dog food")
36+
.put("product_price", 12.99));
37+
38+
if (i%20 ==0) {
39+
System.out.println("Inserted " + i + " documents and count " + vectorStore.count());
40+
}
41+
}
42+
43+
44+
}
45+
46+
47+
@Test
1948
public void quickStartTest() {
20-
String databaseName = "vector_client_test";
21-
String astraToken = System.getenv("ASTRA_DB_APPLICATION_TOKEN");
49+
String databaseName = "vector_client_test";
50+
String vectorStoreName = "demo_store";
51+
String astraToken = System.getenv("ASTRA_DB_APPLICATION_TOKEN");
2252

53+
// 1a. Initialization with a client
2354
AstraVectorClient vectorClient = new AstraVectorClient(astraToken);
55+
56+
// 1b. Create DB (Skip if you already have a database running)
2457
if (!vectorClient.isDatabaseExists(databaseName)) {
2558
vectorClient.createDatabase(databaseName);
2659
}
2760

28-
// Without ODM Accessing the Vector DB with JSON-ISH
61+
// 2. Create a store (delete if exist)
2962
AstraVectorDatabaseClient vectorDb = vectorClient.database(databaseName);
30-
JsonVectorStore jsonVectorStore =
31-
vectorDb.createVectorStore("demo_product", 14);
32-
33-
// ======== INSERTIONS =========
34-
35-
jsonVectorStore.insert(new JsonDocument("doc1")
63+
vectorDb.deleteStore(vectorStoreName);
64+
vectorDb.createVectorStore(vectorStoreName, 14);
65+
66+
// 3. Insert data in the store
67+
JsonVectorStore vectorStore = vectorDb.vectorStore(vectorStoreName);
68+
// 3a. Insert One (attributes as key/value)
69+
vectorStore.insert(new JsonDocument()
70+
.id("doc1") // generated if not set
71+
.vector(new float[]{1f, 0f, 1f, 1f, 1f, 1f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f})
3672
.put("product_name", "HealthyFresh - Beef raw dog food")
37-
.put("product_price", 12.99)
38-
.vector(new float[]{1f, 0f, 1f, 1f, 1f, 1f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f}));
39-
40-
jsonVectorStore.insert(new JsonDocument("doc2")
73+
.put("product_price", 12.99));
74+
// 3b. Insert One (attributes as JSON)
75+
vectorStore.insert(new JsonDocument()
76+
.id("doc2")
77+
.vector(new float[]{1f, 1f, 1f, 1f, 1f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f})
78+
.data("{"
79+
+" \"product_name\": \"HealthyFresh - Chicken raw dog food\", "
80+
+ " \"product_price\": 9.99"
81+
+ "}")
82+
);
83+
// 3c. Insert One (attributes as a MAP)
84+
vectorStore.insert(new JsonDocument()
85+
.id("doc3")
86+
.vector(new float[]{1f, 1f, 1f, 1f, 1f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f})
4187
.data(Map.of("product_name", "HealthyFresh - Chicken raw dog food"))
42-
.vector(new float[]{1f, 1f, 1f, 1f, 1f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f}));
43-
44-
jsonVectorStore.insert(new JsonDocument("doc3")
45-
.data("{"
46-
+" \"product_name\": \"HealthyFresh - Chicken raw dog food\", "
47-
+ " \"product_price\": 9.99, "
48-
+ "}")
49-
.vector(new float[]{1f, 1f, 1f, 1f, 1f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f}));
50-
51-
//jsonVectorStore.insert("{"
52-
// + " \"_id\":\"doc4\","
53-
// + " \"$vector\":[1f, 1f, 1f, 1f, 1f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f],"
54-
// + " \"product_name\": \"HealthyFresh - Chicken raw dog food\", "
55-
// + " \"product_price\": 9.99, "
56-
// + "}");
88+
);
89+
// 3d. Insert as a single Big JSON
90+
vectorStore.insert("{"
91+
+ " \"_id\":\"doc4\","
92+
+ " \"$vector\":[1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],"
93+
+ " \"product_name\": \"HealthyFresh - Chicken raw dog food\", "
94+
+ " \"product_price\": 9.99"
95+
+ "}");
5796

5897
// With ODM
59-
VectorStore<Product> productVectorStore =
60-
vectorDb.createVectorStore("demo_product", 14, Product.class);
98+
VectorStore<Product> productVectorStore = vectorDb.vectorStore(vectorStoreName, Product.class);
6199

62100
// 3 fields: id, payload, vector
63101
productVectorStore.insert("doc5",
@@ -70,6 +108,7 @@ public void quickStartTest() {
70108
new float[]{1f, 1f, 1f, 1f, 1f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f, 0f});
71109
productVectorStore.insert(doc6);
72110

111+
Assertions.assertEquals(6, productVectorStore.count());
73112
}
74113

75114

‎astra-sdk-vector/src/test/resources/logback-test.xml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
<appender-ref ref="STDOUT" />
1717
</logger>
1818

19-
<logger name="io.stargate.sdk.json" level="DEBUG" additivity="false">
19+
<logger name="io.stargate.sdk.json" level="INFO" additivity="false">
2020
<appender-ref ref="STDOUT" />
2121
</logger>
2222

0 commit comments

Comments
 (0)