490 lines
14 KiB
C#
490 lines
14 KiB
C#
// <copyright file="ChatResponseStreamer.cs" company="StellaOps">
|
|
// Copyright (c) StellaOps. Licensed under the BUSL-1.1.
|
|
// </copyright>
|
|
|
|
|
|
using Microsoft.Extensions.Logging;
|
|
using System.Collections.Immutable;
|
|
using System.Globalization;
|
|
using System.Runtime.CompilerServices;
|
|
using System.Text;
|
|
using System.Text.Json;
|
|
|
|
namespace StellaOps.AdvisoryAI.Chat;
|
|
|
|
/// <summary>
|
|
/// Streams chat responses as Server-Sent Events.
|
|
/// Sprint: SPRINT_20260107_006_003 Task CH-006
|
|
/// </summary>
|
|
public sealed class ChatResponseStreamer
|
|
{
|
|
private readonly ILogger<ChatResponseStreamer> _logger;
|
|
private readonly StreamingOptions _options;
|
|
|
|
/// <summary>
|
|
/// Initializes a new instance of the <see cref="ChatResponseStreamer"/> class.
|
|
/// </summary>
|
|
public ChatResponseStreamer(
|
|
ILogger<ChatResponseStreamer> logger,
|
|
StreamingOptions? options = null)
|
|
{
|
|
_logger = logger;
|
|
_options = options ?? new StreamingOptions();
|
|
}
|
|
|
|
/// <summary>
|
|
/// Streams response tokens from an LLM as Server-Sent Events.
|
|
/// </summary>
|
|
/// <param name="tokenSource">The source of tokens from the LLM.</param>
|
|
/// <param name="conversationId">The conversation ID.</param>
|
|
/// <param name="turnId">The turn ID being generated.</param>
|
|
/// <param name="cancellationToken">Cancellation token.</param>
|
|
/// <returns>Async enumerable of SSE events.</returns>
|
|
public async IAsyncEnumerable<StreamEvent> StreamResponseAsync(
|
|
IAsyncEnumerable<TokenChunk> tokenSource,
|
|
string conversationId,
|
|
string turnId,
|
|
[EnumeratorCancellation] CancellationToken cancellationToken = default)
|
|
{
|
|
var contentBuilder = new StringBuilder();
|
|
var citations = new List<CitationEvent>();
|
|
var actions = new List<ActionEvent>();
|
|
var tokenCount = 0;
|
|
var startTime = DateTimeOffset.UtcNow;
|
|
|
|
// Send start event
|
|
yield return new StreamEvent(StreamEventType.Start, new StartEventData
|
|
{
|
|
ConversationId = conversationId,
|
|
TurnId = turnId,
|
|
Timestamp = startTime.ToString("O", CultureInfo.InvariantCulture)
|
|
});
|
|
|
|
await foreach (var chunk in tokenSource.WithCancellation(cancellationToken).ConfigureAwait(false))
|
|
{
|
|
tokenCount++;
|
|
contentBuilder.Append(chunk.Content);
|
|
|
|
// Yield token event
|
|
yield return new StreamEvent(StreamEventType.Token, new TokenEventData
|
|
{
|
|
Content = chunk.Content,
|
|
Index = tokenCount
|
|
});
|
|
|
|
// Check for citations in the accumulated content
|
|
var newCitations = ExtractNewCitations(contentBuilder.ToString(), citations.Count);
|
|
foreach (var citation in newCitations)
|
|
{
|
|
citations.Add(citation);
|
|
yield return new StreamEvent(StreamEventType.Citation, citation);
|
|
}
|
|
|
|
// Check for action proposals
|
|
var newActions = ExtractNewActions(contentBuilder.ToString(), actions.Count);
|
|
foreach (var action in newActions)
|
|
{
|
|
actions.Add(action);
|
|
yield return new StreamEvent(StreamEventType.Action, action);
|
|
}
|
|
|
|
// Periodically send progress events
|
|
if (tokenCount % _options.ProgressInterval == 0)
|
|
{
|
|
yield return new StreamEvent(StreamEventType.Progress, new ProgressEventData
|
|
{
|
|
TokensGenerated = tokenCount,
|
|
ElapsedMs = (int)(DateTimeOffset.UtcNow - startTime).TotalMilliseconds
|
|
});
|
|
}
|
|
}
|
|
|
|
// Send completion event
|
|
var endTime = DateTimeOffset.UtcNow;
|
|
var groundingScore = CalculateGroundingScore(citations.Count, contentBuilder.Length);
|
|
|
|
yield return new StreamEvent(StreamEventType.Done, new DoneEventData
|
|
{
|
|
TurnId = turnId,
|
|
TotalTokens = tokenCount,
|
|
CitationCount = citations.Count,
|
|
ActionCount = actions.Count,
|
|
GroundingScore = groundingScore,
|
|
DurationMs = (int)(endTime - startTime).TotalMilliseconds,
|
|
Timestamp = endTime.ToString("O", CultureInfo.InvariantCulture)
|
|
});
|
|
|
|
_logger.LogInformation(
|
|
"Stream completed: conversation={ConversationId}, turn={TurnId}, tokens={Tokens}, grounding={Grounding:F2}",
|
|
conversationId, turnId, tokenCount, groundingScore);
|
|
}
|
|
|
|
/// <summary>
|
|
/// Formats a stream event as an SSE string.
|
|
/// </summary>
|
|
public static string FormatAsSSE(StreamEvent evt)
|
|
{
|
|
var sb = new StringBuilder();
|
|
|
|
sb.Append("event: ");
|
|
sb.AppendLine(evt.Type.ToString().ToLowerInvariant());
|
|
|
|
var json = JsonSerializer.Serialize(evt.Data, JsonOptions);
|
|
sb.Append("data: ");
|
|
sb.AppendLine(json);
|
|
|
|
sb.AppendLine(); // Empty line to end the event
|
|
|
|
return sb.ToString();
|
|
}
|
|
|
|
/// <summary>
|
|
/// Handles connection drops by checkpointing.
|
|
/// </summary>
|
|
public StreamCheckpoint CreateCheckpoint(
|
|
string conversationId,
|
|
string turnId,
|
|
int tokenIndex,
|
|
string partialContent)
|
|
{
|
|
return new StreamCheckpoint
|
|
{
|
|
ConversationId = conversationId,
|
|
TurnId = turnId,
|
|
TokenIndex = tokenIndex,
|
|
PartialContent = partialContent,
|
|
CreatedAt = DateTimeOffset.UtcNow
|
|
};
|
|
}
|
|
|
|
/// <summary>
|
|
/// Resumes streaming from a checkpoint.
|
|
/// </summary>
|
|
public async IAsyncEnumerable<StreamEvent> ResumeFromCheckpointAsync(
|
|
StreamCheckpoint checkpoint,
|
|
IAsyncEnumerable<TokenChunk> tokenSource,
|
|
[EnumeratorCancellation] CancellationToken cancellationToken = default)
|
|
{
|
|
// Send resume event
|
|
yield return new StreamEvent(StreamEventType.Resume, new ResumeEventData
|
|
{
|
|
ConversationId = checkpoint.ConversationId,
|
|
TurnId = checkpoint.TurnId,
|
|
ResumedFromToken = checkpoint.TokenIndex
|
|
});
|
|
|
|
// Skip tokens we already have
|
|
var skipCount = checkpoint.TokenIndex;
|
|
var skipped = 0;
|
|
|
|
await foreach (var chunk in tokenSource.WithCancellation(cancellationToken).ConfigureAwait(false))
|
|
{
|
|
if (skipped < skipCount)
|
|
{
|
|
skipped++;
|
|
continue;
|
|
}
|
|
|
|
yield return new StreamEvent(StreamEventType.Token, new TokenEventData
|
|
{
|
|
Content = chunk.Content,
|
|
Index = skipped + 1
|
|
});
|
|
}
|
|
}
|
|
|
|
private List<CitationEvent> ExtractNewCitations(string content, int existingCount)
|
|
{
|
|
var citations = new List<CitationEvent>();
|
|
|
|
// Pattern: [type:path]
|
|
var matches = System.Text.RegularExpressions.Regex.Matches(
|
|
content,
|
|
@"\[(?<type>sbom|reach|runtime|vex|attest|auth|docs):(?<path>[^\]]+)\]");
|
|
|
|
for (int i = existingCount; i < matches.Count; i++)
|
|
{
|
|
var match = matches[i];
|
|
citations.Add(new CitationEvent
|
|
{
|
|
Type = match.Groups["type"].Value,
|
|
Path = match.Groups["path"].Value,
|
|
Index = i + 1,
|
|
Verified = false // Will be verified by GroundingValidator
|
|
});
|
|
}
|
|
|
|
return citations;
|
|
}
|
|
|
|
private List<ActionEvent> ExtractNewActions(string content, int existingCount)
|
|
{
|
|
var actions = new List<ActionEvent>();
|
|
|
|
// Pattern: [Label]{action:type,params}
|
|
var matches = System.Text.RegularExpressions.Regex.Matches(
|
|
content,
|
|
@"\[(?<label>[^\]]+)\]\{action:(?<type>\w+)(?:,(?<params>[^}]*))?\}");
|
|
|
|
for (int i = existingCount; i < matches.Count; i++)
|
|
{
|
|
var match = matches[i];
|
|
actions.Add(new ActionEvent
|
|
{
|
|
Type = match.Groups["type"].Value,
|
|
Label = match.Groups["label"].Value,
|
|
Params = match.Groups["params"].Value,
|
|
Index = i + 1,
|
|
Enabled = true // Will be validated by ActionProposalParser
|
|
});
|
|
}
|
|
|
|
return actions;
|
|
}
|
|
|
|
private static double CalculateGroundingScore(int citationCount, int contentLength)
|
|
{
|
|
if (contentLength == 0)
|
|
{
|
|
return 0;
|
|
}
|
|
|
|
// Rough heuristic: expect ~1 citation per 200 characters
|
|
var expectedCitations = contentLength / 200.0;
|
|
if (expectedCitations < 1)
|
|
{
|
|
expectedCitations = 1;
|
|
}
|
|
|
|
var ratio = citationCount / expectedCitations;
|
|
return Math.Min(1.0, ratio);
|
|
}
|
|
|
|
private static readonly JsonSerializerOptions JsonOptions = new()
|
|
{
|
|
PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
|
|
WriteIndented = false
|
|
};
|
|
}
|
|
|
|
/// <summary>
|
|
/// A token chunk from the LLM.
|
|
/// </summary>
|
|
public sealed record TokenChunk
|
|
{
|
|
/// <summary>Gets the token content.</summary>
|
|
public required string Content { get; init; }
|
|
|
|
/// <summary>Gets optional metadata.</summary>
|
|
public ImmutableDictionary<string, string>? Metadata { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Types of stream events.
|
|
/// </summary>
|
|
public enum StreamEventType
|
|
{
|
|
/// <summary>Stream starting.</summary>
|
|
Start,
|
|
|
|
/// <summary>Token generated.</summary>
|
|
Token,
|
|
|
|
/// <summary>Citation extracted.</summary>
|
|
Citation,
|
|
|
|
/// <summary>Action proposal detected.</summary>
|
|
Action,
|
|
|
|
/// <summary>Progress update.</summary>
|
|
Progress,
|
|
|
|
/// <summary>Stream completed.</summary>
|
|
Done,
|
|
|
|
/// <summary>Error occurred.</summary>
|
|
Error,
|
|
|
|
/// <summary>Stream resumed.</summary>
|
|
Resume
|
|
}
|
|
|
|
/// <summary>
|
|
/// A stream event with type and data.
|
|
/// </summary>
|
|
public sealed record StreamEvent(StreamEventType Type, object Data);
|
|
|
|
/// <summary>
|
|
/// Start event data.
|
|
/// </summary>
|
|
public sealed record StartEventData
|
|
{
|
|
/// <summary>Gets the conversation ID.</summary>
|
|
public required string ConversationId { get; init; }
|
|
|
|
/// <summary>Gets the turn ID.</summary>
|
|
public required string TurnId { get; init; }
|
|
|
|
/// <summary>Gets the timestamp.</summary>
|
|
public required string Timestamp { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Token event data.
|
|
/// </summary>
|
|
public sealed record TokenEventData
|
|
{
|
|
/// <summary>Gets the token content.</summary>
|
|
public required string Content { get; init; }
|
|
|
|
/// <summary>Gets the token index.</summary>
|
|
public required int Index { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Citation event data.
|
|
/// </summary>
|
|
public sealed record CitationEvent
|
|
{
|
|
/// <summary>Gets the citation type.</summary>
|
|
public required string Type { get; init; }
|
|
|
|
/// <summary>Gets the citation path.</summary>
|
|
public required string Path { get; init; }
|
|
|
|
/// <summary>Gets the citation index.</summary>
|
|
public required int Index { get; init; }
|
|
|
|
/// <summary>Gets whether the citation is verified.</summary>
|
|
public bool Verified { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Action event data.
|
|
/// </summary>
|
|
public sealed record ActionEvent
|
|
{
|
|
/// <summary>Gets the action type.</summary>
|
|
public required string Type { get; init; }
|
|
|
|
/// <summary>Gets the action label.</summary>
|
|
public required string Label { get; init; }
|
|
|
|
/// <summary>Gets the action parameters.</summary>
|
|
public required string Params { get; init; }
|
|
|
|
/// <summary>Gets the action index.</summary>
|
|
public required int Index { get; init; }
|
|
|
|
/// <summary>Gets whether the action is enabled.</summary>
|
|
public bool Enabled { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Progress event data.
|
|
/// </summary>
|
|
public sealed record ProgressEventData
|
|
{
|
|
/// <summary>Gets tokens generated so far.</summary>
|
|
public required int TokensGenerated { get; init; }
|
|
|
|
/// <summary>Gets elapsed milliseconds.</summary>
|
|
public required int ElapsedMs { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Done event data.
|
|
/// </summary>
|
|
public sealed record DoneEventData
|
|
{
|
|
/// <summary>Gets the turn ID.</summary>
|
|
public required string TurnId { get; init; }
|
|
|
|
/// <summary>Gets total tokens.</summary>
|
|
public required int TotalTokens { get; init; }
|
|
|
|
/// <summary>Gets citation count.</summary>
|
|
public required int CitationCount { get; init; }
|
|
|
|
/// <summary>Gets action count.</summary>
|
|
public required int ActionCount { get; init; }
|
|
|
|
/// <summary>Gets the grounding score.</summary>
|
|
public required double GroundingScore { get; init; }
|
|
|
|
/// <summary>Gets duration in milliseconds.</summary>
|
|
public required int DurationMs { get; init; }
|
|
|
|
/// <summary>Gets the timestamp.</summary>
|
|
public required string Timestamp { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Error event data.
|
|
/// </summary>
|
|
public sealed record ErrorEventData
|
|
{
|
|
/// <summary>Gets the error code.</summary>
|
|
public required string Code { get; init; }
|
|
|
|
/// <summary>Gets the error message.</summary>
|
|
public required string Message { get; init; }
|
|
|
|
/// <summary>Gets tokens generated before error.</summary>
|
|
public int TokensGenerated { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Resume event data.
|
|
/// </summary>
|
|
public sealed record ResumeEventData
|
|
{
|
|
/// <summary>Gets the conversation ID.</summary>
|
|
public required string ConversationId { get; init; }
|
|
|
|
/// <summary>Gets the turn ID.</summary>
|
|
public required string TurnId { get; init; }
|
|
|
|
/// <summary>Gets the token index resumed from.</summary>
|
|
public required int ResumedFromToken { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Checkpoint for resuming streams.
|
|
/// </summary>
|
|
public sealed record StreamCheckpoint
|
|
{
|
|
/// <summary>Gets the conversation ID.</summary>
|
|
public required string ConversationId { get; init; }
|
|
|
|
/// <summary>Gets the turn ID.</summary>
|
|
public required string TurnId { get; init; }
|
|
|
|
/// <summary>Gets the token index.</summary>
|
|
public required int TokenIndex { get; init; }
|
|
|
|
/// <summary>Gets partial content accumulated.</summary>
|
|
public required string PartialContent { get; init; }
|
|
|
|
/// <summary>Gets when checkpoint was created.</summary>
|
|
public required DateTimeOffset CreatedAt { get; init; }
|
|
}
|
|
|
|
/// <summary>
|
|
/// Options for streaming.
|
|
/// </summary>
|
|
public sealed class StreamingOptions
|
|
{
|
|
/// <summary>
|
|
/// Gets or sets the interval for progress events (in tokens).
|
|
/// Default: 50 tokens.
|
|
/// </summary>
|
|
public int ProgressInterval { get; set; } = 50;
|
|
|
|
/// <summary>
|
|
/// Gets or sets the timeout for idle streams.
|
|
/// Default: 30 seconds.
|
|
/// </summary>
|
|
public TimeSpan IdleTimeout { get; set; } = TimeSpan.FromSeconds(30);
|
|
}
|