CSHARP-5828: Add Rerank stage builder

Author: adelinowona

src/MongoDB.Driver/AggregateFluent.cs

@@ -379,6 +379,15 @@ public override IAggregateFluent<TNewResult> Unwind<TNewResult>(FieldDefinition<
             return WithPipeline(_pipeline.Unwind(field, options));
         }
 
+        public override IAggregateFluent<TResult> Rerank(
+            RerankQuery query,
+            FieldDefinition<TResult> path,
+            int numDocsToRerank,
+            string model)
+        {
+            return WithPipeline(_pipeline.Rerank(query, path, numDocsToRerank, model));
+        }
+
         public override IAggregateFluent<TResult> VectorSearch(
             FieldDefinition<TResult> field,
             QueryVector queryVector,

src/MongoDB.Driver/AggregateFluentBase.cs

@@ -331,6 +331,16 @@ public virtual IAggregateFluent<TNewResult> Unwind<TNewResult>(FieldDefinition<T
             throw new NotImplementedException();
         }
 
+        /// <inheritdoc />
+        public virtual IAggregateFluent<TResult> Rerank(
+            RerankQuery query,
+            FieldDefinition<TResult> path,
+            int numDocsToRerank,
+            string model)
+        {
+            throw new NotImplementedException();
+        }
+
         /// <inheritdoc />
         public virtual IAggregateFluent<TResult> VectorSearch(
             FieldDefinition<TResult> field,

src/MongoDB.Driver/IAggregateFluent.cs

@@ -545,6 +545,20 @@ IAggregateFluent<TResult> UnionWith<TWith>(
         /// <returns>The fluent aggregate interface.</returns>
         IAggregateFluent<TNewResult> Unwind<TNewResult>(FieldDefinition<TResult> field, AggregateUnwindOptions<TNewResult> options = null);
 
+        /// <summary>
+        /// Appends a $rerank stage.
+        /// </summary>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="path">The field to send to the reranker.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <returns>The fluent aggregate interface.</returns>
+        IAggregateFluent<TResult> Rerank(
+            RerankQuery query,
+            FieldDefinition<TResult> path,
+            int numDocsToRerank,
+            string model);
+
         /// <summary>
         /// Appends a vector search stage.
         /// </summary>

src/MongoDB.Driver/IAggregateFluentExtensions.cs

@@ -1041,6 +1041,52 @@ public static IAggregateFluent<TNewResult> Unwind<TResult, TNewResult>(this IAgg
             return IAsyncCursorSourceExtensions.SingleOrDefaultAsync(aggregate.Limit(2), cancellationToken);
         }
 
+        /// <summary>
+        /// Appends a $rerank stage.
+        /// </summary>
+        /// <typeparam name="TResult">The type of the result.</typeparam>
+        /// <typeparam name="TField">The type of the field.</typeparam>
+        /// <param name="aggregate">The aggregate.</param>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="path">The field to send to the reranker.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <returns>The fluent aggregate interface.</returns>
+        public static IAggregateFluent<TResult> Rerank<TResult, TField>(
+            this IAggregateFluent<TResult> aggregate,
+            RerankQuery query,
+            Expression<Func<TResult, TField>> path,
+            int numDocsToRerank,
+            string model)
+        {
+            Ensure.IsNotNull(aggregate, nameof(aggregate));
+
+            return aggregate.Rerank(query, new ExpressionFieldDefinition<TResult>(path), numDocsToRerank, model);
+        }
+
+        /// <summary>
+        /// Appends a $rerank stage.
+        /// </summary>
+        /// <typeparam name="TResult">The type of the result.</typeparam>
+        /// <typeparam name="TField">The type of the field.</typeparam>
+        /// <param name="aggregate">The aggregate.</param>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <param name="paths">The fields to send to the reranker.</param>
+        /// <returns>The fluent aggregate interface.</returns>
+        public static IAggregateFluent<TResult> Rerank<TResult, TField>(
+            this IAggregateFluent<TResult> aggregate,
+            RerankQuery query,
+            int numDocsToRerank,
+            string model,
+            params Expression<Func<TResult, TField>>[] paths)
+        {
+            Ensure.IsNotNull(aggregate, nameof(aggregate));
+
+            return aggregate.AppendStage(PipelineStageDefinitionBuilder.Rerank(query, numDocsToRerank, model, paths));
+        }
+
         /// <summary>
         /// Appends a $vectorSearch stage.
         /// </summary>

src/MongoDB.Driver/PipelineDefinitionBuilder.cs

@@ -1081,6 +1081,104 @@ public static PipelineDefinition<TInput, TOutput> RankFusion<TInput, TIntermedia
             return pipeline.AppendStage(PipelineStageDefinitionBuilder.RankFusion(pipelinesWithWeights, options));
         }
 
+        /// <summary>
+        /// Appends a $rerank stage to the pipeline.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TField">The type of the field.</typeparam>
+        /// <typeparam name="TOutput">The type of the output documents.</typeparam>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="path">The field to send to the reranker.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <returns>A new pipeline with an additional stage.</returns>
+        public static PipelineDefinition<TInput, TOutput> Rerank<TInput, TField, TOutput>(
+            this PipelineDefinition<TInput, TOutput> pipeline,
+            RerankQuery query,
+            Expression<Func<TOutput, TField>> path,
+            int numDocsToRerank,
+            string model)
+        {
+            Ensure.IsNotNull(pipeline, nameof(pipeline));
+            return pipeline.AppendStage(
+                PipelineStageDefinitionBuilder.Rerank(query, path, numDocsToRerank, model),
+                pipeline.OutputSerializer);
+        }
+
+        /// <summary>
+        /// Appends a $rerank stage to the pipeline.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TField">The type of the field.</typeparam>
+        /// <typeparam name="TOutput">The type of the output documents.</typeparam>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <param name="paths">The fields to send to the reranker.</param>
+        /// <returns>A new pipeline with an additional stage.</returns>
+        public static PipelineDefinition<TInput, TOutput> Rerank<TInput, TField, TOutput>(
+            this PipelineDefinition<TInput, TOutput> pipeline,
+            RerankQuery query,
+            int numDocsToRerank,
+            string model,
+            params Expression<Func<TOutput, TField>>[] paths)
+        {
+            Ensure.IsNotNull(pipeline, nameof(pipeline));
+            return pipeline.AppendStage(
+                PipelineStageDefinitionBuilder.Rerank(query, numDocsToRerank, model, paths),
+                pipeline.OutputSerializer);
+        }
+
+        /// <summary>
+        /// Appends a $rerank stage to the pipeline.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TOutput">The type of the output documents.</typeparam>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="path">The field to send to the reranker.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <returns>A new pipeline with an additional stage.</returns>
+        public static PipelineDefinition<TInput, TOutput> Rerank<TInput, TOutput>(
+            this PipelineDefinition<TInput, TOutput> pipeline,
+            RerankQuery query,
+            FieldDefinition<TOutput> path,
+            int numDocsToRerank,
+            string model)
+        {
+            Ensure.IsNotNull(pipeline, nameof(pipeline));
+            return pipeline.AppendStage(
+                PipelineStageDefinitionBuilder.Rerank(query, path, numDocsToRerank, model),
+                pipeline.OutputSerializer);
+        }
+
+        /// <summary>
+        /// Appends a $rerank stage to the pipeline.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TOutput">The type of the output documents.</typeparam>
+        /// <param name="pipeline">The pipeline.</param>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="paths">The fields to send to the reranker.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <returns>A new pipeline with an additional stage.</returns>
+        public static PipelineDefinition<TInput, TOutput> Rerank<TInput, TOutput>(
+            this PipelineDefinition<TInput, TOutput> pipeline,
+            RerankQuery query,
+            IEnumerable<FieldDefinition<TOutput>> paths,
+            int numDocsToRerank,
+            string model)
+        {
+            Ensure.IsNotNull(pipeline, nameof(pipeline));
+            return pipeline.AppendStage(
+                PipelineStageDefinitionBuilder.Rerank(query, paths, numDocsToRerank, model),
+                pipeline.OutputSerializer);
+        }
+
         /// <summary>
         /// Appends a $replaceRoot stage to the pipeline.
         /// </summary>

src/MongoDB.Driver/PipelineStageDefinitionBuilder.cs

@@ -1684,6 +1684,118 @@ public static PipelineStageDefinition<TInput, TOutput> ReplaceWith<TInput, TOutp
             return ReplaceWith(new ExpressionAggregateExpressionDefinition<TInput, TOutput>(newRoot));
         }
 
+        /// <summary>
+        /// Creates a $rerank stage.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TField">The type of the field.</typeparam>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="path">The field to send to the reranker.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <returns>The stage.</returns>
+        public static PipelineStageDefinition<TInput, TInput> Rerank<TInput, TField>(
+            RerankQuery query,
+            Expression<Func<TInput, TField>> path,
+            int numDocsToRerank,
+            string model)
+            => Rerank(
+                query,
+                new ExpressionFieldDefinition<TInput>(path),
+                numDocsToRerank,
+                model);
+
+        /// <summary>
+        /// Creates a $rerank stage.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <typeparam name="TField">The type of the field.</typeparam>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <param name="paths">The fields to send to the reranker.</param>
+        /// <returns>The stage.</returns>
+        public static PipelineStageDefinition<TInput, TInput> Rerank<TInput, TField>(
+            RerankQuery query,
+            int numDocsToRerank,
+            string model,
+            params Expression<Func<TInput, TField>>[] paths)
+            => Rerank(
+                query,
+                paths.Select(p => (FieldDefinition<TInput>)new ExpressionFieldDefinition<TInput>(p)),
+                numDocsToRerank,
+                model);
+
+        /// <summary>
+        /// Creates a $rerank stage.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="path">The field to send to the reranker.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <returns>The stage.</returns>
+        public static PipelineStageDefinition<TInput, TInput> Rerank<TInput>(
+            RerankQuery query,
+            FieldDefinition<TInput> path,
+            int numDocsToRerank,
+            string model)
+        {
+            Ensure.IsNotNull(path, nameof(path));
+            return Rerank(
+                query,
+                [path],
+                numDocsToRerank,
+                model);
+        }
+
+        /// <summary>
+        /// Creates a $rerank stage.
+        /// </summary>
+        /// <typeparam name="TInput">The type of the input documents.</typeparam>
+        /// <param name="query">The rerank query.</param>
+        /// <param name="paths">The fields to send to the reranker.</param>
+        /// <param name="numDocsToRerank">The maximum number of documents to rerank.</param>
+        /// <param name="model">The reranking model name.</param>
+        /// <returns>The stage.</returns>
+        public static PipelineStageDefinition<TInput, TInput> Rerank<TInput>(
+            RerankQuery query,
+            IEnumerable<FieldDefinition<TInput>> paths,
+            int numDocsToRerank,
+            string model)
+        {
+            Ensure.IsNotNull(query, nameof(query));
+            Ensure.IsNotNullOrEmpty(paths, nameof(paths));
+            Ensure.IsGreaterThanZero(numDocsToRerank, nameof(numDocsToRerank));
+            Ensure.IsNotNull(model, nameof(model));
+
+            const string operatorName = "$rerank";
+            var stage = new DelegatedPipelineStageDefinition<TInput, TInput>(
+                operatorName,
+                args =>
+                {
+                    ClientSideProjectionHelper.ThrowIfClientSideProjection(args.DocumentSerializer, operatorName);
+
+                    var renderedPaths = paths.Select(p => p.Render(args).FieldName).ToList();
+                    BsonValue pathValue = renderedPaths.Count == 1
+                        ? renderedPaths[0]
+                        : new BsonArray(renderedPaths);
+
+                    var rerankDocument = new BsonDocument
+                    {
+                        { "query", query.Render() },
+                        { "path", pathValue },
+                        { "numDocsToRerank", numDocsToRerank },
+                        { "model", model }
+                    };
+
+                    var document = new BsonDocument(operatorName, rerankDocument);
+                    return new RenderedPipelineStageDefinition<TInput>(operatorName, document, args.DocumentSerializer);
+                });
+
+            return stage;
+        }
+
         /// <summary>
         /// Creates a $set stage.
         /// </summary>

src/MongoDB.Driver/RerankQuery.cs

@@ -0,0 +1,43 @@
+/* Copyright 2010-present MongoDB Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+using MongoDB.Bson;
+using MongoDB.Driver.Core.Misc;
+
+namespace MongoDB.Driver
+{
+    /// <summary>
+    /// Represents a query for the $rerank aggregation stage.
+    /// </summary>
+    public sealed class RerankQuery
+    {
+        private readonly BsonDocument _rendered;
+
+        private RerankQuery(BsonDocument rendered)
+        {
+            _rendered = rendered;
+        }
+
+        /// <summary>
+        /// Creates a text-based rerank query.
+        /// </summary>
+        /// <param name="text">The text to rerank against.</param>
+        /// <returns>A text rerank query.</returns>
+        public static RerankQuery Text(string text) =>
+            new RerankQuery(new BsonDocument("text", Ensure.IsNotNull(text, nameof(text))));
+
+        internal BsonDocument Render() => _rendered;
+    }
+}