Implement Response Streaming

clients/dotnet/WebClient/MemoryWebClient.cs

@@ -11,6 +11,7 @@
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Runtime.CompilerServices;
 using Microsoft.KernelMemory.Context;
 using Microsoft.KernelMemory.Internals;
 
@@ -371,6 +372,60 @@ public async Task<MemoryAnswer> AskAsync(
         return JsonSerializer.Deserialize<MemoryAnswer>(json, s_caseInsensitiveJsonOptions) ?? new MemoryAnswer();
     }
 
+    /// <inheritdoc />
+    public async IAsyncEnumerable<MemoryAnswer> AskAsyncChunk(
+        string question,
+        string? index = null,
+        MemoryFilter? filter = null,
+        ICollection<MemoryFilter>? filters = null,
+        double minRelevance = 0,
+        IContext? context = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        if (filter != null)
+        {
+            if (filters == null) { filters = new List<MemoryFilter>(); }
+
+            filters.Add(filter);
+        }
+
+        MemoryQuery request = new()
+        {
+            Index = index,
+            Question = question,
+            Filters = (filters is { Count: > 0 }) ? filters.ToList() : new(),
+            MinRelevance = minRelevance,
+            ContextArguments = (context?.Arguments ?? new Dictionary<string, object?>()).ToDictionary(),
+        };
+        using StringContent content = new(JsonSerializer.Serialize(request), Encoding.UTF8, "application/json");
+
+        var url = Constants.HttpAskChunkEndpoint.CleanUrlPath();
+        HttpResponseMessage response = await this._client.PostAsync(url, content, cancellationToken).ConfigureAwait(false);
+        response.EnsureSuccessStatusCode();
+        var stream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
+        var eosToken = context.GetCustomEosTokenOrDefault("#DONE#");
+        using (var reader = new StreamReader(stream, Encoding.UTF8))
+        {
+            string? line;
+            while ((line = await reader.ReadLineAsync(cancellationToken).ConfigureAwait(false)) != null)
+            {
+                if (line.StartsWith("data:", StringComparison.Ordinal))
+                {
+                    var jsonData = line.Substring(6);
+                    var chunk = JsonSerializer.Deserialize<MemoryAnswer>(jsonData);
+                    if (chunk != null)
+                    {
+                        yield return chunk;
+                        if (chunk.Result == eosToken)
+                        {
+                            break;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     #region private
 
     private static (string contentType, long contentLength, DateTimeOffset lastModified) GetFileDetails(HttpResponseMessage response)

service/Abstractions/Constants.cs

@@ -61,6 +61,9 @@ public static class TextGeneration
         public static class Rag
 
         {
+            // Used to override EosToken config for streaming
+            public const string EosToken = "custom_rag_eos_token_str";
+
             // Used to override No Answer config
             public const string EmptyAnswer = "custom_rag_empty_answer_str";
 
@@ -128,6 +131,8 @@ public static class Summary
     public const string ReservedPayloadVectorGeneratorField = "vector_generator";
 
     // Endpoints
+
+    public const string HttpAskChunkEndpoint = "/ask_stream";
     public const string HttpAskEndpoint = "/ask";
     public const string HttpSearchEndpoint = "/search";
     public const string HttpDownloadEndpoint = "/download";

service/Abstractions/Context/IContext.cs

@@ -100,6 +100,16 @@ public static string GetCustomEmptyAnswerTextOrDefault(this IContext? context, s
         return defaultValue;
     }
 
+    public static string GetCustomEosTokenOrDefault(this IContext? context, string defaultValue)
+    {
+        if (context.TryGetArg<string>(Constants.CustomContext.Rag.EosToken, out var customValue))
+        {
+            return customValue;
+        }
+
+        return defaultValue;
+    }
+
     public static string GetCustomRagFactTemplateOrDefault(this IContext? context, string defaultValue)
     {
         if (context.TryGetArg<string>(Constants.CustomContext.Rag.FactTemplate, out var customValue))

service/Abstractions/IKernelMemory.cs

@@ -228,4 +228,13 @@ public Task<MemoryAnswer> AskAsync(
         double minRelevance = 0,
         IContext? context = null,
         CancellationToken cancellationToken = default);
+
+    public IAsyncEnumerable<MemoryAnswer> AskAsyncChunk(
+        string question,
+        string? index = null,
+        MemoryFilter? filter = null,
+        ICollection<MemoryFilter>? filters = null,
+        double minRelevance = 0,
+        IContext? context = null,
+        CancellationToken cancellationToken = default);
 }

service/Abstractions/Search/ISearchClient.cs

@@ -50,6 +50,14 @@ Task<MemoryAnswer> AskAsync(
         IContext? context = null,
         CancellationToken cancellationToken = default);
 
+    IAsyncEnumerable<MemoryAnswer> AskAsyncChunk(
+        string index,
+        string question,
+        ICollection<MemoryFilter>? filters = null,
+        double minRelevance = 0,
+        IContext? context = null,
+        CancellationToken cancellationToken = default);
+
     /// <summary>
     /// List the available memory indexes (aka collections).
     /// </summary>

service/Core/MemoryServerless.cs

@@ -288,4 +288,30 @@ public Task<MemoryAnswer> AskAsync(
             context: context,
             cancellationToken: cancellationToken);
     }
+
+    public IAsyncEnumerable<MemoryAnswer> AskAsyncChunk(
+        string question,
+        string? index = null,
+        MemoryFilter? filter = null,
+        ICollection<MemoryFilter>? filters = null,
+        double minRelevance = 0,
+        IContext? context = null,
+        CancellationToken cancellationToken = default)
+    {
+        if (filter != null)
+        {
+            if (filters == null) { filters = new List<MemoryFilter>(); }
+
+            filters.Add(filter);
+        }
+
+        index = IndexName.CleanName(index, this._defaultIndexName);
+        return this._searchClient.AskAsyncChunk(
+            index: index,
+            question: question,
+            filters: filters,
+            minRelevance: minRelevance,
+            context: context,
+            cancellationToken: cancellationToken);
+    }
 }

service/Core/MemoryService.cs

@@ -253,4 +253,33 @@ public Task<MemoryAnswer> AskAsync(
             context: context,
             cancellationToken: cancellationToken);
     }
+
+    public IAsyncEnumerable<MemoryAnswer> AskAsyncChunk(
+        string question,
+        string? index = null,
+        MemoryFilter? filter = null,
+        ICollection<MemoryFilter>? filters = null,
+        double minRelevance = 0,
+        IContext? context = null,
+        CancellationToken cancellationToken = default)
+    {
+        if (filter != null)
+        {
+            if (filters == null)
+            {
+                filters = new List<MemoryFilter>();
+            }
+
+            filters.Add(filter);
+        }
+
+        index = IndexName.CleanName(index, this._defaultIndexName);
+        return this._searchClient.AskAsyncChunk(
+            index: index,
+            question: question,
+            filters: filters,
+            minRelevance: minRelevance,
+            context: context,
+            cancellationToken: cancellationToken);
+    }
 }

service/Core/Search/AnswerGenerator.cs

@@ -95,35 +95,41 @@ internal async Task<MemoryAnswer> GenerateAnswerAsync(
             this._log.LogTrace("Answer generated in {0} msecs", watch.ElapsedMilliseconds);
         }
 
-        // Validate the LLM output
+        return await this.ModeratedAnswerAsync(result.AskResult, cancellationToken).ConfigureAwait(false);
+    }
+
+    internal async Task<MemoryAnswer> ModeratedAnswerAsync(MemoryAnswer currentResult, CancellationToken cancellationToken)
+    {
         if (this._contentModeration != null && this._config.UseContentModeration)
         {
-            var isSafe = await this._contentModeration.IsSafeAsync(result.AskResult.Result, cancellationToken).ConfigureAwait(false);
+            var isSafe = await this._contentModeration.IsSafeAsync(currentResult.Result, cancellationToken).ConfigureAwait(false);
             if (!isSafe)
             {
                 this._log.LogWarning("Unsafe answer detected. Returning error message instead.");
-                this._log.LogSensitive("Unsafe answer: {0}", result.AskResult.Result);
-                result.AskResult.NoResultReason = "Content moderation failure";
-                result.AskResult.Result = this._config.ModeratedAnswer;
+                this._log.LogSensitive("Unsafe answer: {0}", currentResult.Result);
+                currentResult.NoResultReason = "Content moderation failure";
+                currentResult.Result = this._config.ModeratedAnswer;
+                currentResult.NoResult = true;
             }
         }
 
-        return result.AskResult;
+        return currentResult;
     }
 
-    private IAsyncEnumerable<string> GenerateAnswerTokensAsync(string question, string facts, IContext? context, CancellationToken cancellationToken)
+    internal IAsyncEnumerable<string> GenerateAnswerTokensAsync(string question, string facts, IContext? context, CancellationToken cancellationToken)
     {
         string prompt = context.GetCustomRagPromptOrDefault(this._answerPrompt);
         int maxTokens = context.GetCustomRagMaxTokensOrDefault(this._config.AnswerTokens);
         double temperature = context.GetCustomRagTemperatureOrDefault(this._config.Temperature);
         double nucleusSampling = context.GetCustomRagNucleusSamplingOrDefault(this._config.TopP);
-
+        string emptyAnswer = context.GetCustomEmptyAnswerTextOrDefault(this._config.EmptyAnswer);
         prompt = prompt.Replace("{{$facts}}", facts.Trim(), StringComparison.OrdinalIgnoreCase);
 
         question = question.Trim();
         question = question.EndsWith('?') ? question : $"{question}?";
         prompt = prompt.Replace("{{$input}}", question, StringComparison.OrdinalIgnoreCase);
-        prompt = prompt.Replace("{{$notFound}}", this._config.EmptyAnswer, StringComparison.OrdinalIgnoreCase);
+        prompt = prompt.Replace("{{$notFound}}", emptyAnswer, StringComparison.OrdinalIgnoreCase);
+        this._log.LogInformation("New prompt: {0}", prompt);
 
         var options = new TextGenerationOptions
         {

service/Core/Search/SearchClient.cs

@@ -5,6 +5,8 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Globalization;
 using System.Linq;
+using System.Runtime.CompilerServices;
+using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
@@ -174,6 +176,125 @@ public async Task<MemoryAnswer> AskAsync(
         return await this._answerGenerator.GenerateAnswerAsync(question, result, context, cancellationToken).ConfigureAwait(false);
     }
 
+    public async IAsyncEnumerable<MemoryAnswer> AskAsyncChunk(
+        string index,
+        string question,
+        ICollection<MemoryFilter>? filters = null,
+        double minRelevance = 0,
+        IContext? context = null,
+        [EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        this._log.LogInformation("question: '{0}'", question);
+        string emptyAnswer = context.GetCustomEmptyAnswerTextOrDefault(this._config.EmptyAnswer);
+        string eosToken = context.GetCustomEosTokenOrDefault("#DONE#");
+        string answerPrompt = context.GetCustomRagPromptOrDefault(this._answerPrompt);
+        int limit = context.GetCustomRagMaxMatchesCountOrDefault(this._config.MaxMatchesCount);
+        var maxTokens = this._config.MaxAskPromptSize > 0
+            ? this._config.MaxAskPromptSize
+            : this._textGenerator.MaxTokenTotal;
+        SearchClientResult result = SearchClientResult.AskResultInstance(
+            question: question,
+            emptyAnswer: emptyAnswer,
+            maxGroundingFacts: limit,
+            tokensAvailable: maxTokens
+                             - this._textGenerator.CountTokens(answerPrompt)
+                             - this._textGenerator.CountTokens(question)
+                             - this._config.AnswerTokens
+        );
+        if (string.IsNullOrEmpty(question))
+        {
+            this._log.LogWarning("No question provided");
+            yield return result.AskResult;
+            result.AskResult.Result = eosToken;
+            this._log.LogInformation("Eos token: '{0}", result.AskResult.Result);
+            yield return result.AskResult;
+            yield break;
+        }
+
+        this._log.LogTrace("Fetching relevant memories");
+        IAsyncEnumerable<(MemoryRecord, double)> matches = this._memoryDb.GetSimilarListAsync(
+            index: index,
+            text: question,
+            filters: filters,
+            minRelevance: minRelevance,
+            limit: this._config.MaxMatchesCount,
+            withEmbeddings: false,
+            cancellationToken: cancellationToken);
+
+        string factTemplate = context.GetCustomRagFactTemplateOrDefault(this._config.FactTemplate);
+        if (!factTemplate.EndsWith('\n')) { factTemplate += "\n"; }
+
+        // Memories are sorted by relevance, starting from the most relevant
+        await foreach ((MemoryRecord memoryRecord, double recordRelevance) in matches.ConfigureAwait(false))
+        {
+            result.State = SearchState.Continue;
+            result = this.ProcessMemoryRecord(result, index, memoryRecord, recordRelevance, factTemplate);
+
+            if (result.State == SearchState.SkipRecord) { continue; }
+
+            if (result.State == SearchState.Stop) { break; }
+        }
+
+        if (result.FactsAvailableCount > 0 && result.FactsUsedCount == 0)
+        {
+            this._log.LogError("Unable to inject memories in the prompt, not enough tokens available");
+            result.AskResult.NoResultReason = "Unable to use memories";
+            yield return result.AskResult;
+            yield break;
+        }
+
+        if (result.FactsUsedCount == 0)
+        {
+            this._log.LogWarning("No memories available");
+            result.AskResult.NoResultReason = "No memories available";
+            yield return result.AskResult;
+            yield break;
+        }
+
+        this._log.LogTrace("{Count} records processed", result.RecordCount);
+        var charsGenerated = 0;
+        var wholeText = new StringBuilder();
+        await foreach (var x in this._answerGenerator.GenerateAnswerTokensAsync(question, result.Facts.ToString(), context, cancellationToken).ConfigureAwait(true))
+        {
+            var text = new StringBuilder();
+            text.Append(x);
+            wholeText.Append(x);
+            if (this._log.IsEnabled(LogLevel.Trace) && text.Length - charsGenerated >= 30)
+            {
+                charsGenerated = text.Length;
+                this._log.LogTrace("{0} chars generated", charsGenerated);
+            }
+
+            var newAnswer = new MemoryAnswer
+            {
+                Question = question,
+                NoResult = false,
+                Result = text.ToString()
+            };
+            this._log.LogInformation("Chunk: '{0}", newAnswer.Result);
+            yield return newAnswer;
+        }
+
+        var current = new MemoryAnswer
+        {
+            Question = question,
+            NoResult = false,
+            Result = wholeText.ToString(),
+        };
+        MemoryAnswer moderatedResult = await this._answerGenerator.ModeratedAnswerAsync(current, cancellationToken).ConfigureAwait(false);
+        result.AskResult.Result = moderatedResult.Result;
+        result.AskResult.NoResult = moderatedResult.NoResult;
+        result.AskResult.NoResultReason = moderatedResult.NoResultReason;
+        if (result.AskResult.NoResultReason == "Content moderation failure")
+        {
+            yield return result.AskResult;
+        }
+
+        result.AskResult.Result = eosToken;
+        this._log.LogInformation("Eos token: '{0}", result.AskResult.Result);
+        yield return result.AskResult;
+    }
+
     /// <summary>
     /// Process memory records for ASK and SEARCH calls
     /// </summary>
@@ -202,7 +323,7 @@ private SearchClientResult ProcessMemoryRecord(
         // Identify the file in case there are multiple files
         string fileId = record.GetFileId(this._log);
 
-        // Note: this is not a URL and perhaps could be dropped. For now it acts as a unique identifier. See also SourceUrl.
+        // Note: this is not a URL and perhaps could be dropped. For now, it acts as a unique identifier. See also SourceUrl.
         string linkToFile = $"{index}/{documentId}/{fileId}";
 
         // Note: this is "content.url" when importing web pages