Skip to content

Commit c575a1f

Browse files
clunclun
committed
test for BYOV
1 parent e2f2668 commit c575a1f

File tree

14 files changed

+273
-155
lines changed

14 files changed

+273
-155
lines changed

‎astra-db-java/pom.xml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,12 @@
9494
<version>1.0.0-beta2</version>
9595
<scope>test</scope>
9696
</dependency>
97+
<dependency>
98+
<groupId>dev.langchain4j</groupId>
99+
<artifactId>langchain4j-open-ai</artifactId>
100+
<version>1.0.0-beta2</version>
101+
<scope>test</scope>
102+
</dependency>
97103

98104
</dependencies>
99105

‎astra-db-java/src/main/java/com/datastax/astra/client/collections/Collection.java

Lines changed: 37 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -65,13 +65,11 @@
6565
import com.datastax.astra.client.exceptions.DataAPIException;
6666
import com.datastax.astra.client.exceptions.UnexpectedDataAPIResponseException;
6767
import com.datastax.astra.client.tables.commands.options.TableDistinctOptions;
68-
import com.datastax.astra.client.tables.definition.rows.Row;
6968
import com.datastax.astra.internal.api.DataAPIResponse;
7069
import com.datastax.astra.internal.api.DataAPIStatus;
7170
import com.datastax.astra.internal.command.AbstractCommandRunner;
7271
import com.datastax.astra.internal.serdes.DataAPISerializer;
7372
import com.datastax.astra.internal.serdes.collections.DocumentSerializer;
74-
import com.datastax.astra.internal.serdes.tables.RowMapper;
7573
import com.datastax.astra.internal.utils.Assert;
7674
import com.datastax.astra.internal.utils.BetaPreview;
7775
import com.datastax.astra.internal.utils.EscapeUtils;
@@ -103,7 +101,6 @@
103101
import static com.datastax.astra.client.core.options.DataAPIClientOptions.MAX_COUNT;
104102
import static com.datastax.astra.client.exceptions.DataAPIException.ERROR_CODE_INTERRUPTED;
105103
import static com.datastax.astra.client.exceptions.DataAPIException.ERROR_CODE_TIMEOUT;
106-
import static com.datastax.astra.internal.serdes.tables.RowMapper.mapFromRow;
107104
import static com.datastax.astra.internal.utils.AnsiUtils.cyan;
108105
import static com.datastax.astra.internal.utils.AnsiUtils.green;
109106
import static com.datastax.astra.internal.utils.AnsiUtils.magenta;
@@ -793,8 +790,8 @@ private Callable<CollectionInsertManyResult> getInsertManyResultCallable(List<?
793790
Command insertMany = new Command("insertMany")
794791
.withDocuments(documents.subList(start, end))
795792
.withOptions(new Document()
796-
.append(INPUT_ORDERED, collectionInsertManyOptions.isOrdered())
797-
.append(INPUT_RETURN_DOCUMENT_RESPONSES, collectionInsertManyOptions.isReturnDocumentResponses()));
793+
.append(OPTIONS_ORDERED, collectionInsertManyOptions.isOrdered())
794+
.append(OPTIONS_RETURN_DOCUMENT_RESPONSES, collectionInsertManyOptions.isReturnDocumentResponses()));
798795

799796
DataAPIStatus status = runCommand(insertMany, collectionInsertManyOptions).getStatus();
800797
CollectionInsertManyResult result = new CollectionInsertManyResult();
@@ -894,8 +891,8 @@ public Optional<T> findOne(Filter filter, CollectionFindOneOptions findOneOption
894891
.withSort(findOneOptions.getSortArray())
895892
.withProjection(findOneOptions.getProjectionArray())
896893
.withOptions(new Document()
897-
.appendIfNotNull(INPUT_INCLUDE_SIMILARITY, findOneOptions.includeSimilarity())
898-
.appendIfNotNull(INPUT_INCLUDE_SORT_VECTOR, findOneOptions.includeSortVector())
894+
.appendIfNotNull(OPTIONS_INCLUDE_SIMILARITY, findOneOptions.includeSimilarity())
895+
.appendIfNotNull(OPTIONS_INCLUDE_SORT_VECTOR, findOneOptions.includeSortVector())
899896
);
900897

901898
return Optional
@@ -1073,6 +1070,19 @@ public Page<T> findPage(Filter filter, CollectionFindOptions options) {
10731070
// --- Find and Rerank ----
10741071
// -----------------------------
10751072

1073+
/**
1074+
* Finds all documents in the collection.
1075+
*
1076+
* @param options
1077+
* options of find one
1078+
* @return
1079+
* the find iterable interface
1080+
*/
1081+
@BetaPreview
1082+
public CollectionFindAndRerankCursor<T,T> findAndRerank(CollectionFindAndRerankOptions options) {
1083+
return findAndRerank(null, options, getDocumentClass());
1084+
}
1085+
10761086
/**
10771087
* Finds all documents in the collection.
10781088
*
@@ -1093,6 +1103,7 @@ public <R> CollectionFindAndRerankCursor<T, R> findAndRerank(Filter filter, Coll
10931103
return new CollectionFindAndRerankCursor<>(this, filter, options, newRowType);
10941104
}
10951105

1106+
10961107
@BetaPreview
10971108
public <R> Page<RerankResult<R>> findAndRerankPage(Filter filter, CollectionFindAndRerankOptions options, Class<R> newRowType) {
10981109
Command findAndRerankCommand = Command
@@ -1103,12 +1114,13 @@ public <R> Page<RerankResult<R>> findAndRerankPage(Filter filter, CollectionFind
11031114
.withSort(options.getSortArray())
11041115
.withProjection(options.getProjectionArray())
11051116
.withOptions(new Document()
1106-
.appendIfNotNull("rerankOn", options.rerankOn())
1107-
.appendIfNotNull("limit", options.limit())
1108-
.appendIfNotNull("hybridLimits", options.hybridLimits())
1109-
.appendIfNotNull(INPUT_INCLUDE_SORT_VECTOR, options.includeSortVector())
1110-
.appendIfNotNull(INPUT_INCLUDE_SCORES, options.includeScores())
1111-
.appendIfNotNull(INPUT_INCLUDE_SIMILARITY, options.includeSimilarity()));
1117+
.appendIfNotNull(OPTIONS_RERANK_QUERY, options.rerankQuery())
1118+
.appendIfNotNull(OPTIONS_RERANK_ON, options.rerankOn())
1119+
.appendIfNotNull(OPTIONS_LIMIT, options.limit())
1120+
.appendIfNotNull(OPTIONS_HYBRID_LIMITS, options.hybridLimits())
1121+
.appendIfNotNull(OPTIONS_INCLUDE_SORT_VECTOR, options.includeSortVector())
1122+
.appendIfNotNull(OPTIONS_INCLUDE_SCORES, options.includeScores())
1123+
);
11121124
}
11131125

11141126
// Responses MOCK for now
@@ -1183,9 +1195,9 @@ public <R> Page<R> findPage(Filter filter, CollectionFindOptions options, Class<
11831195
.withOptions(new Document()
11841196
.appendIfNotNull("skip", options.skip())
11851197
.appendIfNotNull("limit", options.limit())
1186-
.appendIfNotNull(INPUT_PAGE_STATE, options.pageState())
1187-
.appendIfNotNull(INPUT_INCLUDE_SORT_VECTOR, options.includeSortVector())
1188-
.appendIfNotNull(INPUT_INCLUDE_SIMILARITY, options.includeSimilarity()));
1198+
.appendIfNotNull(OPTIONS_PAGE_STATE, options.pageState())
1199+
.appendIfNotNull(OPTIONS_INCLUDE_SORT_VECTOR, options.includeSortVector())
1200+
.appendIfNotNull(OPTIONS_INCLUDE_SIMILARITY, options.includeSimilarity()));
11891201
}
11901202
DataAPIResponse apiResponse = runCommand(findCommand, options);
11911203

@@ -1619,8 +1631,8 @@ public Optional<T> findOneAndReplace(Filter filter, T replacement, CollectionFin
16191631
.withSort(options.getSortArray())
16201632
.withProjection(options.getProjectionArray())
16211633
.withOptions(new Document()
1622-
.appendIfNotNull(INPUT_UPSERT, options.upsert())
1623-
.appendIfNotNull(INPUT_RETURN_DOCUMENT, options.returnDocument())
1634+
.appendIfNotNull(OPTIONS_UPSERT, options.upsert())
1635+
.appendIfNotNull(OPTIONS_RETURN_DOCUMENT, options.returnDocument())
16241636
);
16251637

16261638
DataAPIResponse res = runCommand(findOneAndReplace, options);
@@ -1667,8 +1679,8 @@ public CollectionUpdateResult replaceOne(Filter filter, T replacement, Collectio
16671679
.withFilter(filter)
16681680
.withReplacement(replacement)
16691681
.withOptions(new Document()
1670-
.appendIfNotNull(INPUT_UPSERT, collectionReplaceOneOptions.upsert())
1671-
.append(INPUT_RETURN_DOCUMENT, ReturnDocument.BEFORE.getKey())
1682+
.appendIfNotNull(OPTIONS_UPSERT, collectionReplaceOneOptions.upsert())
1683+
.append(OPTIONS_RETURN_DOCUMENT, ReturnDocument.BEFORE.getKey())
16721684
);
16731685

16741686
// Execute the `findOneAndReplace`
@@ -1761,8 +1773,8 @@ public Optional<T> findOneAndUpdate(Filter filter, Update update, CollectionFind
17611773
.withSort(options.getSortArray())
17621774
.withProjection(options.getProjectionArray())
17631775
.withOptions(new Document()
1764-
.appendIfNotNull(INPUT_UPSERT, options.upsert())
1765-
.append(INPUT_RETURN_DOCUMENT, options.returnDocument())
1776+
.appendIfNotNull(OPTIONS_UPSERT, options.upsert())
1777+
.append(OPTIONS_RETURN_DOCUMENT, options.returnDocument())
17661778
);
17671779

17681780
DataAPIResponse res = runCommand(cmd, options);
@@ -1810,7 +1822,7 @@ public CollectionUpdateResult updateOne(Filter filter, Update update, Collection
18101822
.withUpdate(update)
18111823
.withSort(updateOptions.getSortArray())
18121824
.withOptions(new Document()
1813-
.appendIfNotNull(INPUT_UPSERT, updateOptions.upsert())
1825+
.appendIfNotNull(OPTIONS_UPSERT, updateOptions.upsert())
18141826
);
18151827
return getUpdateResult(runCommand(cmd, updateOptions));
18161828
}
@@ -1879,8 +1891,8 @@ public CollectionUpdateResult updateMany(Filter filter, Update update, Collectio
18791891
.withFilter(filter)
18801892
.withUpdate(update)
18811893
.withOptions(new Document()
1882-
.appendIfNotNull(INPUT_UPSERT, options.upsert())
1883-
.appendIfNotNull(INPUT_PAGE_STATE, nextPageState));
1894+
.appendIfNotNull(OPTIONS_UPSERT, options.upsert())
1895+
.appendIfNotNull(OPTIONS_PAGE_STATE, nextPageState));
18841896
DataAPIResponse res = runCommand(cmd, options);
18851897
// Data
18861898
if (res.getData() != null) {

‎astra-db-java/src/main/java/com/datastax/astra/client/collections/commands/options/CollectionFindAndRerankOptions.java

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,11 @@ public class CollectionFindAndRerankOptions extends BaseOptions<CollectionFindAn
6666
*/
6767
String rerankOn;
6868

69+
/**
70+
* Rerank Query.
71+
*/
72+
String rerankQuery;
73+
6974
/**
7075
* Options for hybrid projection
7176
*/
@@ -76,8 +81,6 @@ public class CollectionFindAndRerankOptions extends BaseOptions<CollectionFindAn
7681
*/
7782
Boolean includeSortVector;
7883

79-
Boolean includeSimilarity;
80-
8184
/**
8285
* Default constructor.
8386
*/
@@ -187,4 +190,8 @@ public CollectionFindAndRerankOptions rerankOn(String rerankOn) {
187190
}
188191

189192

193+
public CollectionFindAndRerankOptions rerankQuery(String $lexical) {
194+
this.rerankQuery = rerankOn;
195+
return this;
196+
}
190197
}

‎astra-db-java/src/main/java/com/datastax/astra/client/collections/definition/documents/Document.java

Lines changed: 5 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,8 @@
4545
import com.datastax.astra.internal.serdes.DataAPISerializer;
4646
import com.datastax.astra.internal.serdes.collections.DocumentSerializer;
4747
import com.datastax.astra.internal.utils.Assert;
48-
import com.datastax.astra.internal.utils.EscapeUtils;
4948
import com.datastax.astra.internal.utils.BetaPreview;
49+
import com.datastax.astra.internal.utils.EscapeUtils;
5050
import com.fasterxml.jackson.annotation.JsonAnyGetter;
5151
import com.fasterxml.jackson.annotation.JsonAnySetter;
5252
import com.fasterxml.jackson.annotation.JsonIgnore;
@@ -391,29 +391,14 @@ public Document vectorize(String passage, String setPassage) {
391391
/**
392392
* Add a hybrid attribute in the documentW
393393
*
394-
* @param passage
395-
* value for the vectorize attribute
396-
* @return
397-
* self reference
398-
*/
399-
@BetaPreview
400-
public Document hybrid(String passage) {
401-
return appendIfNotNull(DataAPIKeywords.HYBRID.getKeyword(), new Hybrid(passage));
402-
}
403-
404-
/**
405-
* Add a vectorize attribute to the document.
406-
*
407-
* @param vectorize
408-
* string to converted to vector
409-
* @param lexical
410-
* string to be used for lexical search
394+
* @param hybrid
395+
* value for hybrid
411396
* @return
412397
* self reference
413398
*/
414399
@BetaPreview
415-
public Document hybrid(String vectorize, String lexical) {
416-
return appendIfNotNull(DataAPIKeywords.HYBRID.getKeyword(), new Hybrid(vectorize, lexical));
400+
public Document hybrid(Hybrid hybrid) {
401+
return appendIfNotNull(DataAPIKeywords.HYBRID.getKeyword(), hybrid);
417402
}
418403

419404
/**

‎astra-db-java/src/main/java/com/datastax/astra/client/core/hybrid/Hybrid.java

Lines changed: 107 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -23,37 +23,131 @@
2323
import com.datastax.astra.client.collections.definition.documents.Document;
2424
import com.datastax.astra.client.core.DataAPIKeywords;
2525
import com.datastax.astra.client.core.lexical.Lexical;
26+
import com.datastax.astra.client.core.vector.DataAPIVector;
2627
import com.datastax.astra.client.core.vectorize.Vectorize;
2728
import com.datastax.astra.internal.serdes.core.HybridSerializer;
2829
import com.fasterxml.jackson.databind.annotation.JsonSerialize;
2930
import lombok.Data;
3031

32+
/**
33+
* Hybrid object that can be used to store both vector and lexical information.
34+
*/
3135
@Data
3236
@JsonSerialize(using = HybridSerializer.class)
3337
public class Hybrid {
3438

39+
/**
40+
* Document to use.
41+
*/
3542
final Document doc;
3643

44+
/**
45+
* Default constructor.
46+
*/
47+
public Hybrid() {
48+
this.doc = new Document();
49+
}
50+
51+
/**
52+
* Constructor with text.
53+
*
54+
* @param text
55+
* text to use
56+
*/
3757
public Hybrid(String text) {
38-
this.doc = new Document()
39-
.append(DataAPIKeywords.VECTORIZE.getKeyword(), text)
40-
.append(DataAPIKeywords.LEXICAL.getKeyword(), text);
58+
this();
59+
vectorize(text);
60+
lexical(text);
4161
}
4262

43-
public Hybrid(String vectorize, String lexical) {
44-
this.doc = new Document()
45-
.append(DataAPIKeywords.VECTORIZE.getKeyword(), vectorize)
46-
.append(DataAPIKeywords.LEXICAL.getKeyword(), lexical);
63+
/**
64+
* Constructor with custom document.
65+
*
66+
* @param doc
67+
* document to use
68+
*/
69+
public Hybrid(Document doc) {
70+
this.doc = doc;
4771
}
4872

49-
public Hybrid(Vectorize vectorize, Lexical lexical) {
50-
this.doc = new Document()
51-
.append(DataAPIKeywords.VECTORIZE.getKeyword(), vectorize)
52-
.append(DataAPIKeywords.LEXICAL.getKeyword(), lexical);
73+
/**
74+
* Add a vectorize field.
75+
*
76+
* @param vectorize
77+
* vectorize to use
78+
* @return
79+
* this
80+
*/
81+
public Hybrid vectorize(String vectorize) {
82+
this.doc.append(DataAPIKeywords.VECTORIZE.getKeyword(), vectorize);
83+
return this;
5384
}
5485

55-
public Hybrid(Document doc) {
56-
this.doc = doc;
86+
/**
87+
* Add a vectorize field.
88+
*
89+
* @param vectorize
90+
* vectorize to use
91+
* @return
92+
* this
93+
*/
94+
public Hybrid vectorize(Vectorize vectorize) {
95+
this.doc.append(DataAPIKeywords.VECTORIZE.getKeyword(), vectorize);
96+
return this;
97+
}
98+
99+
/**
100+
* Add a lexical field.
101+
*
102+
* @param lexical
103+
* lexical to use
104+
* @return
105+
* this
106+
*/
107+
public Hybrid lexical(String lexical) {
108+
this.doc.append(DataAPIKeywords.LEXICAL.getKeyword(), lexical);
109+
return this;
57110
}
58111

112+
/**
113+
* Add a lexical field.
114+
*
115+
* @param lexical
116+
* lexical to use
117+
* @return
118+
* this
119+
*/
120+
public Hybrid lexical(Lexical lexical) {
121+
this.doc.append(DataAPIKeywords.LEXICAL.getKeyword(), lexical);
122+
return this;
123+
}
124+
125+
/**
126+
* Add a vector that will be serialized as a float array.
127+
*
128+
* @param embeddings
129+
* embeddings to use
130+
* @return
131+
* this
132+
*/
133+
public Hybrid vector(float[] embeddings) {
134+
this.doc.append(DataAPIKeywords.VECTOR.getKeyword(), embeddings);
135+
return this;
136+
}
137+
138+
/**
139+
* Add a vector that will be serialized as a base64 encoded string..
140+
*
141+
* @param vector
142+
* vector to use
143+
* @return
144+
* this
145+
*/
146+
public Hybrid vector(DataAPIVector vector) {
147+
this.doc.append(DataAPIKeywords.VECTOR.getKeyword(), vector);
148+
return this;
149+
}
150+
151+
152+
59153
}

0 commit comments

Comments
 (0)