Update find and Rerank, start reworking the tests

datastax · clun · Apr 4, 2025 · Nov 21, 2024 · Nov 30, 2024 · Nov 30, 2024
commit e2f26686316feddbeefcc747a9a186a42f52c0cc
diff --git a/TEST.MD b/TEST.MD
@@ -0,0 +1,52 @@
+
+## Run against Local HCD Instance 
+(_to run locally your must have the data api running locally on op of HCD or DSE_)
+
+```bash
+mvn clean test -Ptest_local
+```
+
+> **Why it works ?**
+> 
+> - In `pom.xml` the environment variable `ASTRA_DB_JAVA_TEST_ENV` is set to `local` 
+> - Tests are annotated with
+> 
+> ```java
+> @EnabledIfSystemProperty(named = "ASTRA_DB_JAVA_TEST_ENV", matches = "local")
+> @DisabledIfSystemProperty(named = "ASTRA_DB_JAVA_TEST_ENV", matches = "(?!local)")
+> public class MyTest {
+> }
+> ```
+
+- To run with IDE make sure to set the environment variable `ASTRA_DB_JAVA_TEST_ENV` to `local` in your Run configuration
+
+
+## Run against Astra DEV
+
+- `ASTRA_DB_APPLICATION_TOKEN_DEV` should be defined in your environment variables, it is not in the different POM for privacy reasons.
+
+- To run with Maven (AWS,`eu-west-2`)
+
+```bash
+mvn clean test -Ptest_astra_dev
+```
+
+- To run with Maven on One particular region
+
+```java
+mvn clean test -Ptest_astra_dev \
+  -DASTRA_CLOUD_PROVIDER_DEV=GCP \
+  -DASTRA_CLOUD_REGION_DEV=us-central-1
+```
+
+- To run overriding also the token
+
+```java
+mvn clean test \
+  -Ptest_astra_dev \
+  -DASTRA_CLOUD_PROVIDER_DEV=GCP \
+  -DASTRA_CLOUD_REGION_DEV=us-central-1 \
+  -ASTRA_DB_APPLICATION_TOKEN_DEV=AstraCS:... \
+  -Dtest="com.datastax.astra.test.integration.*.*Test"
+```
+
diff --git a/astra-db-java-tools/src/test/java/com/datastax/astra/samples/CsvPhilosophers.java b/astra-db-java-tools/src/test/java/com/datastax/astra/samples/CsvPhilosophers.java
@@ -30,6 +30,7 @@ public static void main(String[] args) throws Exception {
         CsvLoader.load(csvFilename, collection, new CsvRowMapper() {
             @Override
             public Document map(Document csvRow) {
+
                 // Tags should be an Array
                 csvRow.vectorize(csvRow.getString("quote"));
                 csvRow.append("tags",csvRow.getString("tags").split(";"));

diff --git a/astra-db-java/pom.xml b/astra-db-java/pom.xml
@@ -86,6 +86,15 @@
             <artifactId>awaitility</artifactId>
             <scope>test</scope>
         </dependency>
+
+        <!-- Embedding model in memory for testing -->
+        <dependency>
+            <groupId>dev.langchain4j</groupId>
+            <artifactId>langchain4j-embeddings-all-minilm-l6-v2</artifactId>
+            <version>1.0.0-beta2</version>
+            <scope>test</scope>
+        </dependency>
+
     </dependencies>
 
     <build>

diff --git a/astra-db-java/src/main/java/com/datastax/astra/client/admin/AstraDBDatabaseAdmin.java b/astra-db-java/src/main/java/com/datastax/astra/client/admin/AstraDBDatabaseAdmin.java
@@ -25,6 +25,7 @@
 import com.datastax.astra.client.core.options.DataAPIClientOptions;
 import com.datastax.astra.client.databases.commands.results.FindEmbeddingProvidersResult;
 import com.datastax.astra.client.databases.DatabaseOptions;
+import com.datastax.astra.client.databases.commands.results.FindRerankingProvidersResult;
 import com.datastax.astra.internal.api.AstraApiEndpoint;
 import com.datastax.astra.internal.command.AbstractCommandRunner;
 import com.datastax.astra.internal.utils.Assert;
@@ -192,6 +193,14 @@ public FindEmbeddingProvidersResult findEmbeddingProviders() {
         return new FindEmbeddingProvidersResult(admin.findEmbeddingProviders().getEmbeddingProviders());
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public FindRerankingProvidersResult findRerankingProviders() {
+        log.debug("findRerankingProviders");
+        DataAPIDatabaseAdmin admin = new DataAPIDatabaseAdmin(db, this.options);
+        return new FindRerankingProvidersResult(admin.findRerankingProviders().getRerankingProviders());
+    }
+
     /** {@inheritDoc} */
     @Override
     public void createKeyspace(String keyspace, boolean updateDBKeyspace) {

diff --git a/astra-db-java/src/main/java/com/datastax/astra/client/admin/DataAPIDatabaseAdmin.java b/astra-db-java/src/main/java/com/datastax/astra/client/admin/DataAPIDatabaseAdmin.java
@@ -23,9 +23,11 @@
 import com.datastax.astra.client.core.options.BaseOptions;
 import com.datastax.astra.client.core.commands.Command;
 import com.datastax.astra.client.core.commands.CommandType;
+import com.datastax.astra.client.core.rerank.RerankProvider;
 import com.datastax.astra.client.databases.commands.results.FindEmbeddingProvidersResult;
 import com.datastax.astra.client.core.vectorize.EmbeddingProvider;
 import com.datastax.astra.client.databases.Database;
+import com.datastax.astra.client.databases.commands.results.FindRerankingProvidersResult;
 import com.datastax.astra.client.databases.definition.keyspaces.KeyspaceOptions;
 import com.datastax.astra.internal.api.DataAPIResponse;
 import com.datastax.astra.internal.command.AbstractCommandRunner;
@@ -109,6 +111,14 @@ public FindEmbeddingProvidersResult findEmbeddingProviders() {
                 EmbeddingProvider.class));
     }
 
+    @Override
+    public FindRerankingProvidersResult findRerankingProviders() {
+        DataAPIResponse res = runCommand(Command.create("findRerankingProviders"));
+        return new FindRerankingProvidersResult(
+                res.getStatusKeyAsMap("rerankingProviders",
+                        RerankProvider.class));
+    }
+
     /** {@inheritDoc} */
     @Override
     public Database getDatabase() {

diff --git a/astra-db-java/src/main/java/com/datastax/astra/client/admin/DatabaseAdmin.java b/astra-db-java/src/main/java/com/datastax/astra/client/admin/DatabaseAdmin.java
@@ -21,10 +21,12 @@
  */
 
 import com.datastax.astra.client.core.options.BaseOptions;
+import com.datastax.astra.client.core.rerank.RerankProvider;
 import com.datastax.astra.client.databases.Database;
 import com.datastax.astra.client.core.commands.CommandRunner;
 import com.datastax.astra.client.core.vectorize.EmbeddingProvider;
 import com.datastax.astra.client.databases.commands.results.FindEmbeddingProvidersResult;
+import com.datastax.astra.client.databases.commands.results.FindRerankingProvidersResult;
 import com.datastax.astra.internal.utils.Assert;
 
 import java.util.Set;
@@ -81,6 +83,24 @@ public interface DatabaseAdmin {
      */
     FindEmbeddingProvidersResult findEmbeddingProviders();
 
+    /**
+     * Retrieve the list of reranking providers available in the current database. Reranking providers are services
+     * that sort a list of record based on a algorithm (eg bm25) . This method returns a map of provider names to
+     * {@link RerankProvider} instances, allowing applications
+     * to access and utilize the reranking services.
+     *
+     * <p>Example usage:</p>
+     * <pre>
+     * {@code
+     * // Assuming 'client' is an instance of DataApiClient
+     * Map<String, EmbeddingProvider> providers = client.findEmbeddingProvidersAsMap());
+     * }
+     * </pre>
+     * @return
+     *      list of available providers
+     */
+    FindRerankingProvidersResult findRerankingProviders();
+
     /**
      * Asynchronously retrieves a stream of keyspaces names available in the current database. This method facilitates
      * non-blocking operations by allowing the application to continue executing other tasks while the list of keyspace

diff --git a/astra-db-java/src/main/java/com/datastax/astra/client/collections/Collection.java b/astra-db-java/src/main/java/com/datastax/astra/client/collections/Collection.java
@@ -59,7 +59,7 @@
 import com.datastax.astra.client.core.query.Filter;
 import com.datastax.astra.client.core.query.Filters;
 import com.datastax.astra.client.core.query.Projection;
-import com.datastax.astra.client.core.reranking.RerankResult;
+import com.datastax.astra.client.core.rerank.RerankResult;
 import com.datastax.astra.client.core.vector.DataAPIVector;
 import com.datastax.astra.client.databases.Database;
 import com.datastax.astra.client.exceptions.DataAPIException;
@@ -73,6 +73,7 @@
 import com.datastax.astra.internal.serdes.collections.DocumentSerializer;
 import com.datastax.astra.internal.serdes.tables.RowMapper;
 import com.datastax.astra.internal.utils.Assert;
+import com.datastax.astra.internal.utils.BetaPreview;
 import com.datastax.astra.internal.utils.EscapeUtils;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
@@ -1082,31 +1083,32 @@ public Page<T> findPage(Filter filter, CollectionFindOptions options) {
      * @return
      *      the find iterable interface
      */
+    @BetaPreview
     public CollectionFindAndRerankCursor<T,T> findAndRerank(Filter filter, CollectionFindAndRerankOptions options) {
         return findAndRerank(filter, options, getDocumentClass());
     }
 
+    @BetaPreview
     public <R> CollectionFindAndRerankCursor<T, R> findAndRerank(Filter filter, CollectionFindAndRerankOptions options, Class<R> newRowType) {
         return new CollectionFindAndRerankCursor<>(this, filter, options, newRowType);
     }
 
+    @BetaPreview
     public <R> Page<RerankResult<R>> findAndRerankPage(Filter filter, CollectionFindAndRerankOptions options, Class<R> newRowType) {
         Command findAndRerankCommand = Command
                 .create("findAndRerank")
                 .withFilter(filter);
         if (options != null) {
             findAndRerankCommand
-                    .withSort(options.getSortArray())
-                    .withProjection(options.getProjectionArray())
-                    .withOptions(new Document()
-                            .appendIfNotNull("rerankOn", options.rerankOn())
-                            .appendIfNotNull("limit", options.limit())
-                            .appendIfNotNull("hybridProjection", options.hybridProjection().getValue())
-                            .appendIfNotNull("hybridLimits", options.hybridLimits())
-                            .appendIfNotNull(INPUT_INCLUDE_SORT_VECTOR, options.includeSortVector())
-                            .appendIfNotNull(INPUT_INCLUDE_SIMILARITY, options.includeSimilarity())
-                    )
-            ;
+              .withSort(options.getSortArray())
+              .withProjection(options.getProjectionArray())
+              .withOptions(new Document()
+                  .appendIfNotNull("rerankOn", options.rerankOn())
+                  .appendIfNotNull("limit", options.limit())
+                  .appendIfNotNull("hybridLimits", options.hybridLimits())
+                  .appendIfNotNull(INPUT_INCLUDE_SORT_VECTOR, options.includeSortVector())
+                  .appendIfNotNull(INPUT_INCLUDE_SCORES, options.includeScores())
+                  .appendIfNotNull(INPUT_INCLUDE_SIMILARITY, options.includeSimilarity()));
         }
 
         // Responses MOCK for now
@@ -1139,10 +1141,6 @@ public <R> Page<RerankResult<R>> findAndRerankPage(Filter filter, CollectionFind
             DocumentSerializer serializer = new DocumentSerializer();
             R results1 = serializer.convertValue(document, newRowType);
 
-            // MAP WITH ROW FUNCTION
-            Row row = RowMapper.mapAsRow(document);
-            R result = RowMapper.mapFromRow(row, getSerializer(), newRowType);
-
             // Getting associated document response
             Document documentResponse = documentResponses.get(i);
             Map<String, Double> scores = documentResponse.getMap("scores", String.class, Double.class);
@@ -1199,12 +1197,7 @@ public <R> Page<R> findPage(Filter filter, CollectionFindOptions options, Class<
         return new Page<>(
                 apiResponse.getData().getNextPageState(),
                 apiResponse.getData().getDocuments().stream()
-                        .map(d -> {
-                            Row row = RowMapper.mapAsRow(d);
-                            return mapFromRow(row, getSerializer(), newRowType);
-                        })
-                        // .map(d -> d.map(newRowType))
-                        //.map(d -> RowMapper.mapFromRow(d, getSerializer(), newRowType))
+                        .map(d -> d.map(newRowType))
                         .collect(Collectors.toList()), sortVector);
     }
 

diff --git a/.../com/datastax/astra/client/collections/commands/cursor/CollectionFindAndRerankCursor.java b/.../com/datastax/astra/client/collections/commands/cursor/CollectionFindAndRerankCursor.java
@@ -26,7 +26,7 @@
 import com.datastax.astra.client.core.query.Filter;
 import com.datastax.astra.client.core.query.Projection;
 import com.datastax.astra.client.core.query.Sort;
-import com.datastax.astra.client.core.reranking.RerankResult;
+import com.datastax.astra.client.core.rerank.RerankResult;
 import com.datastax.astra.internal.command.AbstractCursor;
 import lombok.Getter;
 

diff --git a/...main/java/com/datastax/astra/client/collections/commands/cursor/CollectionFindCursor.java b/...main/java/com/datastax/astra/client/collections/commands/cursor/CollectionFindCursor.java
@@ -22,6 +22,7 @@
 
 import com.datastax.astra.client.collections.Collection;
 import com.datastax.astra.client.collections.commands.options.CollectionFindOptions;
+import com.datastax.astra.client.collections.definition.documents.Document;
 import com.datastax.astra.client.core.paging.CursorState;
 import com.datastax.astra.client.core.query.Filter;
 import com.datastax.astra.client.core.query.Projection;

diff --git a/...om/datastax/astra/client/collections/commands/options/CollectionFindAndRerankOptions.java b/...om/datastax/astra/client/collections/commands/options/CollectionFindAndRerankOptions.java
@@ -23,7 +23,6 @@
 import com.datastax.astra.client.core.commands.CommandType;
 import com.datastax.astra.client.core.hybrid.Hybrid;
 import com.datastax.astra.client.core.hybrid.HybridLimits;
-import com.datastax.astra.client.core.hybrid.HybridProjection;
 import com.datastax.astra.client.core.options.BaseOptions;
 import com.datastax.astra.client.core.query.Projection;
 import com.datastax.astra.client.core.query.Sort;
@@ -70,7 +69,7 @@ public class CollectionFindAndRerankOptions extends BaseOptions<CollectionFindAn
     /**
      * Options for hybrid projection
      */
-    HybridProjection hybridProjection;
+    Boolean includeScores;
 
     /**
      * Flag to include sortVector in the result when operating a semantic search.
@@ -187,15 +186,5 @@ public CollectionFindAndRerankOptions rerankOn(String rerankOn) {
         return this;
     }
 
-    /**
-     * Add a hybridProjection clause in the find block
-     *
-     * @param hybridProjection value for hybridProjection options
-     * @return current command
-     */
-    public CollectionFindAndRerankOptions hybridProjection(HybridProjection hybridProjection) {
-        this.hybridProjection = hybridProjection;
-        return this;
-    }
 
 }
diff --git a/.../src/main/java/com/datastax/astra/client/collections/definition/CollectionDefinition.java b/.../src/main/java/com/datastax/astra/client/collections/definition/CollectionDefinition.java
@@ -23,8 +23,8 @@
 import com.datastax.astra.client.core.lexical.Analyzer;
 import com.datastax.astra.client.core.lexical.AnalyzerTypes;
 import com.datastax.astra.client.core.lexical.LexicalOptions;
-import com.datastax.astra.client.core.reranking.RerankingOptions;
-import com.datastax.astra.client.core.reranking.RerankingServiceOptions;
+import com.datastax.astra.client.core.rerank.CollectionRerankOptions;
+import com.datastax.astra.client.core.rerank.RerankServiceOptions;
 import com.datastax.astra.client.core.vector.SimilarityMetric;
 import com.datastax.astra.client.core.vector.VectorOptions;
 import com.datastax.astra.client.core.vectorize.VectorServiceOptions;
@@ -346,6 +346,13 @@ public LexicalOptions getLexical() {
         return lexical;
     }
 
+    public CollectionDefinition disableLexical() {
+        if (getLexical() == null) {
+            lexical(new LexicalOptions().enabled(false));
+        }
+        return this;
+    }
+
     /**
      * Builder pattern.
      *
@@ -393,27 +400,27 @@ public CollectionDefinition lexical(AnalyzerTypes analyzerType) {
     /**
      * Reranking options
      */
-    private RerankingOptions reranking;
+    private CollectionRerankOptions rerank;
 
     /**
      * Gets reranking
      *
      * @return value of reranking
      */
-    public RerankingOptions getReranking() {
-        return reranking;
+    public CollectionRerankOptions getRerank() {
+        return rerank;
     }
 
     /**
      * Builder pattern.
      *
-     * @param rerankingOptions
+     * @param collectionRerankOptions
      *      reranking service information
      * @return
      *      self reference
      */
-    public CollectionDefinition reranking(RerankingOptions rerankingOptions) {
-        reranking = rerankingOptions;
+    public CollectionDefinition rerank(CollectionRerankOptions collectionRerankOptions) {
+        rerank = collectionRerankOptions;
         return this;
     }
 
@@ -424,13 +431,20 @@ public CollectionDefinition reranking(RerankingOptions rerankingOptions) {
      *      reranker
      * @return self reference
      */
-    public CollectionDefinition reranking(String reranker) {
-        if (getReranking() == null) {
-            reranking = new RerankingOptions().enabled(true);
+    public CollectionDefinition rerank(String reranker) {
+        if (getRerank() == null) {
+            rerank = new CollectionRerankOptions().enabled(true);
         }
-        getReranking()
+        getRerank()
                 .enabled(true)
-                .service(new RerankingServiceOptions().provider(reranker));
+                .service(new RerankServiceOptions().provider(reranker));
+        return this;
+    }
+
+    public CollectionDefinition disableRerank() {
+        if (getRerank() == null) {
+            rerank = new CollectionRerankOptions().enabled(false);
+        }
         return this;
     }