| | | 1 | | using Microsoft.Extensions.AI; |
| | | 2 | | |
| | | 3 | | namespace NexusLabs.Needlr.AgentFramework.Diagnostics; |
| | | 4 | | |
| | | 5 | | /// <summary> |
| | | 6 | | /// Default implementation of <see cref="IDagRunResult"/>. |
| | | 7 | | /// </summary> |
| | | 8 | | [DoNotAutoRegister] |
| | | 9 | | internal sealed class DagRunResult : IDagRunResult |
| | | 10 | | { |
| | | 11 | | private readonly Lazy<IReadOnlyDictionary<string, ChatResponse?>> _lazyResponses; |
| | | 12 | | private readonly Lazy<TokenUsage?> _lazyAggregateTokenUsage; |
| | | 13 | | |
| | 39 | 14 | | internal DagRunResult( |
| | 39 | 15 | | IReadOnlyList<IAgentStageResult> stages, |
| | 39 | 16 | | IReadOnlyDictionary<string, IDagNodeResult> nodeResults, |
| | 39 | 17 | | IReadOnlyDictionary<string, IReadOnlyList<IAgentStageResult>> branchResults, |
| | 39 | 18 | | TimeSpan totalDuration, |
| | 39 | 19 | | bool succeeded, |
| | 39 | 20 | | string? errorMessage, |
| | 39 | 21 | | Exception? exception = null, |
| | 39 | 22 | | int? plannedStageCount = null) |
| | | 23 | | { |
| | 39 | 24 | | Stages = stages; |
| | 39 | 25 | | PlannedStageCount = plannedStageCount ?? stages.Count; |
| | 39 | 26 | | NodeResults = nodeResults; |
| | 39 | 27 | | BranchResults = branchResults; |
| | 39 | 28 | | TotalDuration = totalDuration; |
| | 39 | 29 | | Succeeded = succeeded; |
| | 39 | 30 | | ErrorMessage = errorMessage; |
| | 39 | 31 | | Exception = exception; |
| | | 32 | | |
| | 39 | 33 | | _lazyResponses = new Lazy<IReadOnlyDictionary<string, ChatResponse?>>(() => |
| | 41 | 34 | | stages |
| | 2 | 35 | | .GroupBy(s => s.AgentName) |
| | 45 | 36 | | .ToDictionary(g => g.Key, g => g.Last().FinalResponse)); |
| | | 37 | | |
| | 39 | 38 | | _lazyAggregateTokenUsage = new Lazy<TokenUsage?>(() => |
| | 39 | 39 | | { |
| | 3 | 40 | | var diagnostics = stages |
| | 3 | 41 | | .Select(s => s.Diagnostics) |
| | 3 | 42 | | .Where(d => d is not null) |
| | 3 | 43 | | .ToList(); |
| | 39 | 44 | | |
| | 5 | 45 | | if (diagnostics.Count == 0) return null; |
| | 39 | 46 | | |
| | 1 | 47 | | return new TokenUsage( |
| | 2 | 48 | | InputTokens: diagnostics.Sum(d => d!.AggregateTokenUsage.InputTokens), |
| | 2 | 49 | | OutputTokens: diagnostics.Sum(d => d!.AggregateTokenUsage.OutputTokens), |
| | 2 | 50 | | TotalTokens: diagnostics.Sum(d => d!.AggregateTokenUsage.TotalTokens), |
| | 2 | 51 | | CachedInputTokens: diagnostics.Sum(d => d!.AggregateTokenUsage.CachedInputTokens), |
| | 3 | 52 | | ReasoningTokens: diagnostics.Sum(d => d!.AggregateTokenUsage.ReasoningTokens)); |
| | 39 | 53 | | }); |
| | 39 | 54 | | } |
| | | 55 | | |
| | 2 | 56 | | public IReadOnlyList<IAgentStageResult> Stages { get; } |
| | | 57 | | |
| | 0 | 58 | | public int PlannedStageCount { get; } |
| | | 59 | | |
| | 4 | 60 | | public IReadOnlyDictionary<string, ChatResponse?> FinalResponses => _lazyResponses.Value; |
| | | 61 | | |
| | 3 | 62 | | public TimeSpan TotalDuration { get; } |
| | | 63 | | |
| | 6 | 64 | | public TokenUsage? AggregateTokenUsage => _lazyAggregateTokenUsage.Value; |
| | | 65 | | |
| | 31 | 66 | | public bool Succeeded { get; } |
| | | 67 | | |
| | 31 | 68 | | public string? ErrorMessage { get; } |
| | | 69 | | |
| | 1 | 70 | | public Exception? Exception { get; } |
| | | 71 | | |
| | 11 | 72 | | public IReadOnlyDictionary<string, IDagNodeResult> NodeResults { get; } |
| | | 73 | | |
| | 5 | 74 | | public IReadOnlyDictionary<string, IReadOnlyList<IAgentStageResult>> BranchResults { get; } |
| | | 75 | | } |