| | | 1 | | using System.Text.Json; |
| | | 2 | | using Elsa.Abstractions; |
| | | 3 | | using Elsa.AI.Abstractions.Contracts; |
| | | 4 | | using Elsa.AI.Abstractions.Models; |
| | | 5 | | using Elsa.AI.Host.Endpoints.AI; |
| | | 6 | | using Elsa.AI.Host.Options; |
| | | 7 | | using Elsa.AI.Host.Permissions; |
| | | 8 | | using Elsa.AI.Host.Streaming; |
| | | 9 | | using JetBrains.Annotations; |
| | | 10 | | using Microsoft.AspNetCore.Http; |
| | | 11 | | using Microsoft.Extensions.Options; |
| | | 12 | | |
| | | 13 | | namespace Elsa.AI.Host.Endpoints.AI.Chat; |
| | | 14 | | |
| | | 15 | | [PublicAPI] |
| | 8 | 16 | | public class Endpoint( |
| | 8 | 17 | | IAIOrchestrator orchestrator, |
| | 8 | 18 | | AIStreamSessionManager sessionManager, |
| | 8 | 19 | | IOptions<AIHostOptions> options) : ElsaEndpoint<AIChatRequest> |
| | | 20 | | { |
| | | 21 | | public override void Configure() |
| | | 22 | | { |
| | 0 | 23 | | Post("/ai/chat"); |
| | 0 | 24 | | ConfigurePermissions(AIPermissions.Chat); |
| | 0 | 25 | | } |
| | | 26 | | |
| | | 27 | | public override async Task HandleAsync(AIChatRequest request, CancellationToken cancellationToken) |
| | | 28 | | { |
| | 8 | 29 | | var conversationId = string.IsNullOrWhiteSpace(request.ConversationId) ? Guid.NewGuid().ToString("N") : request. |
| | 8 | 30 | | var userPermissions = AIHttpContextIdentity.GetPermissions(HttpContext); |
| | 8 | 31 | | request = request with |
| | 8 | 32 | | { |
| | 8 | 33 | | ConversationId = conversationId, |
| | 8 | 34 | | Message = request.Message ?? "", |
| | 8 | 35 | | Attachments = request.Attachments ?? [], |
| | 8 | 36 | | IsReconnect = sessionManager.CanReconnect(conversationId), |
| | 8 | 37 | | TenantId = AIHttpContextIdentity.GetTenantId(HttpContext), |
| | 8 | 38 | | UserId = AIHttpContextIdentity.GetActorId(HttpContext), |
| | 8 | 39 | | UserPermissions = userPermissions, |
| | 8 | 40 | | Agent = AIHttpContextIdentity.GetAuthorizedAgent(request.Agent, options.Value, userPermissions), |
| | 8 | 41 | | ProviderName = null |
| | 8 | 42 | | }; |
| | 8 | 43 | | var response = HttpContext.Response; |
| | 8 | 44 | | response.ContentType = "text/event-stream"; |
| | 8 | 45 | | response.Headers["Cache-Control"] = "no-cache"; |
| | | 46 | | |
| | 8 | 47 | | var completed = false; |
| | 8 | 48 | | var reconnectAccepted = request.IsReconnect; |
| | 8 | 49 | | var requestedReconnectConversationId = request.ConversationId; |
| | 8 | 50 | | var reconnectConnected = false; |
| | 8 | 51 | | var disconnectedConversationId = request.ConversationId; |
| | | 52 | | try |
| | | 53 | | { |
| | 32 | 54 | | await foreach (var streamEvent in orchestrator.ExecuteChatAsync(request, cancellationToken)) |
| | | 55 | | { |
| | 8 | 56 | | disconnectedConversationId = streamEvent.ConversationId; |
| | 8 | 57 | | if (reconnectAccepted && !reconnectConnected) |
| | | 58 | | { |
| | 1 | 59 | | if (!string.Equals(requestedReconnectConversationId, disconnectedConversationId, StringComparison.Or |
| | 1 | 60 | | sessionManager.ReleaseReconnect(requestedReconnectConversationId); |
| | | 61 | | |
| | 1 | 62 | | sessionManager.MarkConnected(disconnectedConversationId); |
| | 1 | 63 | | reconnectConnected = true; |
| | | 64 | | } |
| | | 65 | | |
| | 8 | 66 | | await response.WriteAsync($"event: {streamEvent.Type}\n", cancellationToken); |
| | 8 | 67 | | await response.WriteAsync($"data: {JsonSerializer.Serialize(streamEvent)}\n\n", cancellationToken); |
| | 8 | 68 | | await response.Body.FlushAsync(cancellationToken); |
| | 8 | 69 | | } |
| | | 70 | | |
| | 8 | 71 | | completed = true; |
| | 8 | 72 | | } |
| | 0 | 73 | | catch (OperationCanceledException) when (HttpContext.RequestAborted.IsCancellationRequested) |
| | | 74 | | { |
| | | 75 | | // Expected when the client disconnects; the finally block records reconnect state. |
| | 0 | 76 | | return; |
| | | 77 | | } |
| | | 78 | | finally |
| | | 79 | | { |
| | 8 | 80 | | if (!completed) |
| | | 81 | | { |
| | 0 | 82 | | sessionManager.MarkDisconnected(disconnectedConversationId, options.Value.ReconnectGrace); |
| | | 83 | | } |
| | 8 | 84 | | else if (reconnectAccepted && !reconnectConnected) |
| | | 85 | | { |
| | 1 | 86 | | sessionManager.ReleaseReconnect(requestedReconnectConversationId); |
| | | 87 | | } |
| | | 88 | | } |
| | 8 | 89 | | } |
| | | 90 | | } |