feat(agentic): Phase 3 — beam-search decoding (#1544)

- LocalEngineOptions.BeamWidth (>1) switches non-streaming GetResponseAsync to beam search: explores N hypotheses in parallel, expands each by its top tokens (honoring the ITokenConstraint and stop sequences per beam), prunes to the top-N by length-normalized log-probability, and returns the best completion. Deterministic; streaming stays token-by-token.

- LogSoftmax (with allowed-token masking → -inf) + length-normalized scoring + Beam bookkeeping.

- 3 tests (green net10.0 + net471): greedy takes the locally-best token ("A"); beam width 2 discovers the globally better "BC" path the same model would miss greedily; beam respects a finite-state constraint ("AB"). No null-forgiving (removed an introduced allowed! via proper narrowing).

Part of epic #1544 (Phase 3).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: franklinic

src/Agentic/Models/Local/LocalEngineChatClient.cs

@@ -71,13 +71,23 @@ public Task<ChatResponse> GetResponseAsync(
         CancellationToken cancellationToken = default)
     {
         var promptIds = BuildPromptIds(messages);
-        var sampling = ResolveSampling(options);
-        var sampler = new TokenSampler<T>(sampling.Seed);
         var maxTokens = ResolveMaxTokens(options);
-
         var stopSequences = ResolveStopSequences(options);
-        var generated = new List<int>();
-        var finishReason = RunGeneration(promptIds, maxTokens, sampler, sampling, stopSequences, generated, cancellationToken);
+
+        List<int> generated;
+        ChatFinishReason finishReason;
+        var beamWidth = _options.BeamWidth is { } width && width > 1 ? width : 1;
+        if (beamWidth > 1)
+        {
+            (generated, finishReason) = RunBeamSearch(promptIds, maxTokens, beamWidth, stopSequences, cancellationToken);
+        }
+        else
+        {
+            var sampling = ResolveSampling(options);
+            var sampler = new TokenSampler<T>(sampling.Seed);
+            generated = new List<int>();
+            finishReason = RunGeneration(promptIds, maxTokens, sampler, sampling, stopSequences, generated, cancellationToken);
+        }
 
         var text = generated.Count > 0 ? _tokenizer.Decode(generated) : string.Empty;
         text = TrimAtStopSequence(text, stopSequences);
@@ -194,6 +204,152 @@ private ChatFinishReason RunGeneration(
         return sampler.Sample(logits, sampling, allowed);
     }
 
+    // Beam search: explore `beamWidth` hypotheses in parallel, expanding each by its top tokens (honoring
+    // the constraint), and keep the highest-scoring (length-normalized log-probability) beams. Returns the
+    // best completion. Deterministic.
+    private (List<int> Tokens, ChatFinishReason Finish) RunBeamSearch(
+        IReadOnlyList<int> promptIds,
+        int maxTokens,
+        int beamWidth,
+        IReadOnlyList<string>? stopSequences,
+        CancellationToken cancellationToken)
+    {
+        var beams = new List<Beam> { new(new List<int>(), 0.0, finished: false) };
+
+        for (var step = 0; step < maxTokens; step++)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+            if (beams.All(b => b.Finished))
+            {
+                break;
+            }
+
+            var expanded = new List<Beam>();
+            foreach (var beam in beams)
+            {
+                if (beam.Finished)
+                {
+                    expanded.Add(beam);
+                    continue;
+                }
+
+                IReadOnlyCollection<int>? allowed = null;
+                if (_options.Constraint is { } constraint)
+                {
+                    allowed = constraint.AllowedNextTokens(beam.Tokens);
+                    if (allowed is not null && allowed.Count == 0)
+                    {
+                        expanded.Add(new Beam(beam.Tokens, beam.LogProbability, finished: true));
+                        continue;
+                    }
+                }
+
+                var context = new List<int>(promptIds.Count + beam.Tokens.Count);
+                context.AddRange(promptIds);
+                context.AddRange(beam.Tokens);
+
+                var logProbabilities = LogSoftmax(_model.NextTokenLogits(context), allowed);
+                foreach (var (tokenId, logProbability) in TopTokens(logProbabilities, beamWidth))
+                {
+                    var finished = tokenId == _tokenizer.EosTokenId;
+                    var tokens = new List<int>(beam.Tokens);
+                    if (!finished)
+                    {
+                        tokens.Add(tokenId);
+                    }
+
+                    if (!finished && stopSequences is not null && ContainsStopSequence(_tokenizer.Decode(tokens), stopSequences))
+                    {
+                        finished = true;
+                    }
+
+                    expanded.Add(new Beam(tokens, beam.LogProbability + logProbability, finished));
+                }
+            }
+
+            beams = expanded
+                .OrderByDescending(NormalizedScore)
+                .Take(beamWidth)
+                .ToList();
+        }
+
+        var best = beams.OrderByDescending(NormalizedScore).First();
+        return (best.Tokens, best.Finished ? ChatFinishReason.Stop : ChatFinishReason.Length);
+    }
+
+    private static double NormalizedScore(Beam beam) => beam.LogProbability / Math.Max(1, beam.Tokens.Count);
+
+    private static double[] LogSoftmax(Vector<T> logits, IReadOnlyCollection<int>? allowed)
+    {
+        var count = logits.Length;
+        bool[]? allowedMask = null;
+        if (allowed is not null)
+        {
+            allowedMask = new bool[count];
+            foreach (var id in allowed)
+            {
+                if (id >= 0 && id < count)
+                {
+                    allowedMask[id] = true;
+                }
+            }
+        }
+
+        var scores = new double[count];
+        var max = double.NegativeInfinity;
+        for (var i = 0; i < count; i++)
+        {
+            var value = allowedMask is null || allowedMask[i] ? Convert.ToDouble(logits[i]) : double.NegativeInfinity;
+            scores[i] = value;
+            if (value > max)
+            {
+                max = value;
+            }
+        }
+
+        var sumExp = 0.0;
+        for (var i = 0; i < count; i++)
+        {
+            if (!double.IsNegativeInfinity(scores[i]))
+            {
+                sumExp += Math.Exp(scores[i] - max);
+            }
+        }
+
+        var logSumExp = max + Math.Log(sumExp);
+        for (var i = 0; i < count; i++)
+        {
+            scores[i] = double.IsNegativeInfinity(scores[i]) ? double.NegativeInfinity : scores[i] - logSumExp;
+        }
+
+        return scores;
+    }
+
+    private static IEnumerable<(int TokenId, double LogProbability)> TopTokens(double[] logProbabilities, int k)
+    {
+        return Enumerable.Range(0, logProbabilities.Length)
+            .Where(i => !double.IsNegativeInfinity(logProbabilities[i]))
+            .OrderByDescending(i => logProbabilities[i])
+            .Take(k)
+            .Select(i => (i, logProbabilities[i]));
+    }
+
+    private sealed class Beam
+    {
+        public Beam(List<int> tokens, double logProbability, bool finished)
+        {
+            Tokens = tokens;
+            LogProbability = logProbability;
+            Finished = finished;
+        }
+
+        public List<int> Tokens { get; }
+
+        public double LogProbability { get; }
+
+        public bool Finished { get; }
+    }
+
     private static IReadOnlyList<string>? ResolveStopSequences(ChatOptions? options)
     {
         var stops = options?.StopSequences;

src/Agentic/Models/Local/LocalEngineOptions.cs

@@ -40,4 +40,13 @@ public sealed class LocalEngineOptions
     /// vocabulary) at the logits rather than relying on prompting.
     /// </summary>
     public ITokenConstraint? Constraint { get; set; }
+
+    /// <summary>
+    /// Gets or sets the beam width for beam-search decoding. <c>null</c> or a value &lt;= 1 uses ordinary
+    /// token-by-token sampling/greedy decoding. A value &gt; 1 explores that many hypotheses in parallel and
+    /// returns the highest-scoring (length-normalized) completion — deterministic, and typically higher
+    /// quality than greedy for short structured outputs. Beam search applies to non-streaming
+    /// <see cref="LocalEngineChatClient{T}.GetResponseAsync"/>; streaming always decodes token-by-token.
+    /// </summary>
+    public int? BeamWidth { get; set; }
 }

tests/AiDotNet.Tests/UnitTests/Agentic/Local/BeamSearchTests.cs

@@ -0,0 +1,142 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using AiDotNet.Agentic.Models;
+using AiDotNet.Agentic.Models.Local;
+using AiDotNet.LinearAlgebra;
+using Xunit;
+
+namespace AiDotNetTests.UnitTests.Agentic.Local
+{
+    public class BeamSearchTests
+    {
+        // Vocabulary: 0=EOS, 1=A, 2=B, 3=C. The model's next-token logits depend on the last token, set up
+        // so that the locally-best first token (A) leads to a worse overall sequence than B -> C.
+        private static TransitionModel BuildModel() => new(
+            vocabularySize: 4,
+            transitions: new Dictionary<int, double[]>
+            {
+                [0] = new[] { 0.1, 2.0, 1.9, 0.1 }, // start (prompt's last token): A slightly beats B
+                [1] = new[] { 2.0, 0.1, 0.1, 0.1 }, // after A: EOS (A is a short, lower-total path)
+                [2] = new[] { 0.1, 0.1, 0.1, 5.0 }, // after B: C is almost certain
+                [3] = new[] { 5.0, 0.1, 0.1, 0.1 }, // after C: EOS
+            });
+
+        private static MapTokenizer BuildTokenizer() =>
+            new(eos: 0, ("A", 1), ("B", 2), ("C", 3));
+
+        [Fact(Timeout = 60000)]
+        public async Task Greedy_TakesLocallyBestToken()
+        {
+            var engine = new LocalEngineChatClient<double>(BuildModel(), BuildTokenizer(), options: new LocalEngineOptions
+            {
+                Sampling = new LocalSamplingOptions { Temperature = 0.0 }, // greedy, beam width 1
+            });
+
+            var response = await engine.GetResponseAsync(new[] { ChatMessage.User("go") });
+
+            Assert.Equal("A", response.Text);
+        }
+
+        [Fact(Timeout = 60000)]
+        public async Task BeamSearch_FindsHigherProbabilitySequence()
+        {
+            var engine = new LocalEngineChatClient<double>(BuildModel(), BuildTokenizer(), options: new LocalEngineOptions
+            {
+                BeamWidth = 2,
+            });
+
+            var response = await engine.GetResponseAsync(new[] { ChatMessage.User("go") });
+
+            // Beam search explores B as well as A and discovers the globally better B -> C completion.
+            Assert.Equal("BC", response.Text);
+            Assert.Equal(ChatFinishReason.Stop, response.FinishReason);
+        }
+
+        [Fact(Timeout = 60000)]
+        public async Task BeamSearch_RespectsConstraint()
+        {
+            // Force the grammar A -> B -> (terminal) even though the model would prefer other paths.
+            var constraint = new FiniteStateTokenConstraint(
+                start: new[] { 1 },
+                transitions: new Dictionary<int, IReadOnlyCollection<int>>
+                {
+                    [1] = new[] { 2 },
+                    [2] = Array.Empty<int>(),
+                });
+
+            var engine = new LocalEngineChatClient<double>(BuildModel(), BuildTokenizer(), options: new LocalEngineOptions
+            {
+                BeamWidth = 3,
+                Constraint = constraint,
+            });
+
+            var response = await engine.GetResponseAsync(new[] { ChatMessage.User("go") });
+
+            Assert.Equal("AB", response.Text);
+            Assert.Equal(ChatFinishReason.Stop, response.FinishReason);
+        }
+
+        // ---- Test doubles ----
+
+        private sealed class TransitionModel : ICausalLanguageModel<double>
+        {
+            private readonly Dictionary<int, double[]> _transitions;
+            private readonly double[] _default;
+
+            public TransitionModel(int vocabularySize, Dictionary<int, double[]> transitions)
+            {
+                VocabularySize = vocabularySize;
+                _transitions = transitions;
+                _default = new double[vocabularySize];
+                _default[0] = 5.0; // unknown state -> prefer EOS
+            }
+
+            public int VocabularySize { get; }
+
+            public Vector<double> NextTokenLogits(IReadOnlyList<int> tokenIds)
+            {
+                var last = tokenIds[tokenIds.Count - 1];
+                var logits = _transitions.TryGetValue(last, out var row) ? row : _default;
+                return new Vector<double>((double[])logits.Clone());
+            }
+        }
+
+        private sealed class MapTokenizer : IGenerationTokenizer
+        {
+            private readonly Dictionary<string, int> _wordToId = new(StringComparer.Ordinal);
+            private readonly Dictionary<int, string> _idToWord = new();
+
+            public MapTokenizer(int eos, params (string Word, int Id)[] entries)
+            {
+                EosTokenId = eos;
+                _idToWord[eos] = string.Empty;
+                foreach (var (word, id) in entries)
+                {
+                    _wordToId[word] = id;
+                    _idToWord[id] = word;
+                }
+            }
+
+            public int EosTokenId { get; }
+
+            public IReadOnlyList<int> Encode(string text)
+            {
+                var ids = text
+                    .Split(new[] { ' ', '\n', '\r', '\t' }, StringSplitOptions.RemoveEmptyEntries)
+                    .Select(w => _wordToId.TryGetValue(w, out var id) ? id : EosTokenId)
+                    .ToList();
+                if (ids.Count == 0)
+                {
+                    ids.Add(EosTokenId);
+                }
+
+                return ids;
+            }
+
+            public string Decode(IReadOnlyList<int> tokenIds) =>
+                string.Concat(tokenIds.Select(id => _idToWord.TryGetValue(id, out var w) ? w : string.Empty));
+        }
+    }
+}