SearchEngine.cs

using KeyLoad.Core;
using KeyLoad.Query.Features.Search;
using Microsoft.Extensions.Options;

namespace KeyLoad.Query;

/// <summary>Runs exact text, vector and hybrid ranking within one authorized read cut.</summary>
public sealed partial class SearchEngine
{
    private readonly DatabaseEngine database;
    private readonly ITextProjection? textProjection;
    private readonly QueryExecutionOptions execution;
    private readonly IOptions<QueryExecutionOptions> configuration;

    /// <summary>Creates an authorized search owner with one frozen native execution policy.</summary>
    /// <param name="database">Node-owned canonical database.</param>
    /// <param name="options">Centrally validated query and text execution budgets.</param>
    /// <param name="textProjection">Optional native derived text projection.</param>
    public SearchEngine(DatabaseEngine database, IOptions<QueryExecutionOptions> options, ITextProjection? textProjection = null)
    {
        ArgumentNullException.ThrowIfNull(database);
        ArgumentNullException.ThrowIfNull(options);
        configuration = options;
        execution = options.Value;
        execution.Validate();
        this.database = database;
        this.textProjection = textProjection;
    }

    private const string UnsafeSynchronousSearch = "Await SearchAsync when native text search runs on a scheduler or synchronization context.";

    /// <summary>Returns exact fused ranks with selected documents projected by persisted policy.</summary>
    /// <param name="principalId">Persisted principal identity.</param>
    /// <param name="request">Search fields, query and result limit.</param>
    /// <param name="cancellationToken">Caller cancellation for the complete read.</param>
    /// <returns>Ranked and authorized documents.</returns>
    public RankedDocument[] Search(string principalId, SearchRequest request, CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(request);
        if (request.Text is not null && textProjection is not null
            && (TaskScheduler.Current != TaskScheduler.Default || SynchronizationContext.Current is not null))
        {
            throw new InvalidOperationException(UnsafeSynchronousSearch);
        }
        var budget = new ReadExecutionBudget(database.OperationLimitsOptions, database.EvaluationClock, cancellationToken);
        QueryResultBudgetPolicy.Constrain(budget, execution);
        budget.Check();
        using var reservation = database.AdmitQuery(cancellationToken);
        return SearchCore(principalId, request, budget);
    }

    /// <summary>Awaits exact search and native projection settlement without blocking the caller's execution context.</summary>
    /// <param name="principalId">Persisted principal identity.</param>
    /// <param name="request">Search fields, query and result limit.</param>
    /// <param name="cancellationToken">Caller cancellation for admission and the complete worker.</param>
    /// <returns>Ranked and authorized documents after complete lease settlement.</returns>
    public async Task<RankedDocument[]> SearchAsync(string principalId, SearchRequest request, CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(request);
        var budget = new ReadExecutionBudget(database.OperationLimitsOptions, database.EvaluationClock, cancellationToken);
        QueryResultBudgetPolicy.Constrain(budget, execution);
        budget.Check();
        using var reservation = database.AdmitQuery(cancellationToken);
        return await Task.Run(() => SearchCore(principalId, request, budget), cancellationToken).ConfigureAwait(false);
    }

    /// <summary>Executes graph scope, retrieval and expansion inside one authorized search cut.</summary>
    /// <param name="principalId">Persisted caller identity.</param>
    /// <param name="request">Versioned search and graph operators.</param>
    /// <param name="cancellationToken">Caller cancellation for admission and the complete worker.</param>
    /// <returns>Exact fused hits and optional separately projected graph context.</returns>
    public async Task<GraphSearchResult> GraphSearchAsync(string principalId, GraphSearchRequest request,
        CancellationToken cancellationToken = default)
    {
        ArgumentNullException.ThrowIfNull(request);
        var budget = new ReadExecutionBudget(database.OperationLimitsOptions, database.EvaluationClock, cancellationToken);
        return await GraphSearchAsync(principalId, request, budget, cancellationToken).ConfigureAwait(false);
    }

    internal async Task<GraphSearchResult> GraphSearchAsync(string principalId, GraphSearchRequest request,
        ReadExecutionBudget budget, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(request);
        ArgumentNullException.ThrowIfNull(budget);
        QueryResultBudgetPolicy.Constrain(budget, execution);
        budget.Check();
        using var reservation = database.AdmitQuery(cancellationToken);
        return await Task.Run(() => GraphSearchExecutor.Execute(database, textProjection, principalId, request, budget, execution),
            cancellationToken).ConfigureAwait(false);
    }

    private RankedDocument[] SearchCore(string principalId, SearchRequest request, ReadExecutionBudget budget)
    {
        budget.Check();
        SearchRequestValidation.Validate(request, database.Limits, false, execution);
        FilteredSearchEligibility.ValidateRequest(request.AllowedIds, database.Limits, budget);
        var similarity = request.Vector is { } vector
            ? PreparedSimilarity.Create(vector.AsMemory(), request.Space!.Metric) : null;
        FilteredSearchRequestSizer.EnsureBounded(request, database.Limits.MaxQueryBytes, budget);
        var eligibility = FilteredSearchEligibility.Create(request.AllowedIds, budget);
        var original = database.WithQueryView(principalId, request.Partition, request.Collection,
            (view, principal, resource) => SearchCapturedExecution.Start(database, textProjection, view,
                principal, resource, request, budget, execution, similarity, eligibility));
        return SearchCapturedExecution.Complete(database, original, request, budget, similarity, eligibility);
    }

    /// <summary>Computes the selected vector metric using the same native SIMD blocks and bounded scalar tails as search.</summary>
    /// <param name="left">Finite query vector.</param>
    /// <param name="right">Finite candidate vector.</param>
    /// <param name="metric">Typed distance metric.</param>
    /// <returns>Similarity score; Euclidean distance is negated.</returns>
    public static double Similarity(float[] left, float[] right, DistanceMetric metric)
    {
        ArgumentNullException.ThrowIfNull(left);
        var prepared = PreparedSimilarity.Create(left.AsMemory(), metric);
        ArgumentNullException.ThrowIfNull(right);
        return prepared.Score(right.AsMemory());
    }
}