Merge branch 'feature/chat-span' of private:sdcb/chats into feature/c…

…hat-span
sdcb · Dec 25, 2024 · 9969cd2 · 9969cd2
2 parents 23a7805 + bf7f01c
commit 9969cd2
Show file tree

Hide file tree

Showing 10 changed files with 307 additions and 14 deletions.
diff --git a/src/BE/Controllers/Chats/Chats/ChatController.cs b/src/BE/Controllers/Chats/Chats/ChatController.cs
@@ -268,6 +268,7 @@ public async Task<IActionResult> ChatPrivate(
             chat.Messages.Add(resp.AssistantMessage);
         }
         if (isEmptyChat) chat.Title = request.UserMessage!.Text[..Math.Min(50, request.UserMessage.Text.Length)];
+        chat.LeafMessage = resps.Last().AssistantMessage;
         await db.SaveChangesAsync(cancellationToken);
 
         // yield end messages
@@ -279,6 +280,7 @@ public async Task<IActionResult> ChatPrivate(
         {
             await YieldResponse(SseResponseLine.ResponseMessage(resp.SpanId, resp.AssistantMessage, idEncryption, fup));
         }
+        await YieldResponse(SseResponseLine.ChatLeafMessageId(chat.LeafMessageId!.Value, idEncryption));
 
         // finish costs
         if (resps.Any(x => x.Cost.CostBalance > 0))

diff --git a/src/BE/Controllers/Chats/Chats/Dtos/SseResponseKind.cs b/src/BE/Controllers/Chats/Chats/Dtos/SseResponseKind.cs
@@ -9,4 +9,5 @@ public enum SseResponseKind
     UpdateTitle = 4,
     TitleSegment = 5,
     ResponseMessage = 6,
+    ChatLeafMessageId = 7,
 }
diff --git a/src/BE/Controllers/Chats/Chats/Dtos/SseResponseLine.cs b/src/BE/Controllers/Chats/Chats/Dtos/SseResponseLine.cs
@@ -129,4 +129,13 @@ public static SseResponseLine<string> TitleSegment(string titleSegment)
             Kind = SseResponseKind.TitleSegment,
         };
     }
+
+    public static SseResponseLine<string> ChatLeafMessageId(long leafMessageId, IUrlEncryptionService idEncryption)
+    {
+        return new SseResponseLine<string>
+        {
+            Result = idEncryption.EncryptMessageId(leafMessageId),
+            Kind = SseResponseKind.ChatLeafMessageId,
+        };
+    }
 }
diff --git a/src/BE/Controllers/Chats/UserChats/Dtos/ChatsResponse.cs b/src/BE/Controllers/Chats/UserChats/Dtos/ChatsResponse.cs
@@ -17,8 +17,8 @@ public record ChatsResponse
     [JsonPropertyName("isShared")]
     public required bool IsShared { get; init; }
 
-    [JsonPropertyName("messageCount")]
-    public required int MessageCount { get; init; }
+    [JsonPropertyName("leafMessageId")]
+    public required string? LeafMessageId { get; init; }
 }
 
 public record ChatSpanDto

diff --git a/src/BE/Controllers/Chats/UserChats/Dtos/UpdateChatsRequest.cs b/src/BE/Controllers/Chats/UserChats/Dtos/UpdateChatsRequest.cs
@@ -1,4 +1,6 @@
 using Chats.BE.DB;
+using Chats.BE.Services.UrlEncryption;
+using Microsoft.EntityFrameworkCore;
 using System.Text.Json.Serialization;
 
 namespace Chats.BE.Controllers.Chats.UserChats.Dtos;
@@ -8,15 +10,65 @@ public class UpdateChatsRequest
     [JsonPropertyName("title")]
     public string? Title { get; set; } = null!;
 
-    [JsonPropertyName("modelId")]
-    public short? ModelId { get; set; }
-
     [JsonPropertyName("isShared")]
     public bool? IsShared { get; set; }
 
     [JsonPropertyName("isDeleted")]
     public bool? IsDeleted { get; set; }
 
+    [JsonPropertyName("setsLeafMessageId")]
+    public bool SetsLeafMessageId { get; set; }
+
+    [JsonPropertyName("leafMessageId")]
+    public string? LeafMessageId { get; set; } = null!;
+
+    public DecryptedUpdateChatsRequest Decrypt(IUrlEncryptionService urlEncryptionService)
+    {
+        return new DecryptedUpdateChatsRequest
+        {
+            Title = Title,
+            IsShared = IsShared,
+            IsDeleted = IsDeleted,
+            SetsLeafMessageId = SetsLeafMessageId,
+            LeafMessageId = LeafMessageId switch
+            {
+                null => null,
+                _ => urlEncryptionService.DecryptMessageId(LeafMessageId)
+            },
+        };
+    }
+}
+
+
+public class DecryptedUpdateChatsRequest
+{
+    public string? Title { get; set; } = null!;
+
+    public bool? IsShared { get; set; }
+
+    public bool? IsDeleted { get; set; }
+
+    public bool SetsLeafMessageId { get; set; }
+
+    public long? LeafMessageId { get; set; } = null!;
+
+    public async Task<string?> Validate(ChatsDB db, int chatId)
+    {
+        if (Title != null && Title.Length > 50)
+        {
+            return "Title is too long";
+        }
+
+        if (SetsLeafMessageId && LeafMessageId != null)
+        {
+            if (!await db.Messages.AnyAsync(x => x.Id == LeafMessageId && x.ChatId == chatId))
+            {
+                return "Leaf message not found";
+            }
+        }
+        return null;
+    }
+
     public void ApplyToChats(Chat chat)
     {
         if (Title != null)
@@ -31,5 +83,9 @@ public void ApplyToChats(Chat chat)
         {
             chat.IsDeleted = IsDeleted.Value;
         }
+        if (SetsLeafMessageId)
+        {
+            chat.LeafMessageId = LeafMessageId;
+        }
     }
-}
+}
diff --git a/src/BE/Controllers/Chats/UserChats/UserChatsController.cs b/src/BE/Controllers/Chats/UserChats/UserChatsController.cs
@@ -34,7 +34,7 @@ public async Task<ActionResult<ChatsResponse>> GetOneChat(string encryptedChatId
                     Temperature = s.Temperature,
                     EnableSearch = s.EnableSearch,
                 }).ToArray(),
-                MessageCount = x.Messages.Count,
+                LeafMessageId = x.LeafMessageId != null ? idEncryption.EncryptMessageId(x.LeafMessageId.Value) : null,
             })
             .FirstOrDefaultAsync(cancellationToken);
 
@@ -72,7 +72,7 @@ public async Task<ActionResult<PagedResult<ChatsResponse>>> GetChats([FromQuery]
                     Temperature = s.Temperature,
                     EnableSearch = s.EnableSearch,
                 }).ToArray(),
-                MessageCount = x.Messages.Count,
+                LeafMessageId = x.LeafMessageId != null ? idEncryption.EncryptMessageId(x.LeafMessageId.Value) : null,
             }),
             request,
             cancellationToken);
@@ -144,7 +144,7 @@ public async Task<ActionResult<ChatsResponse>> CreateChat([FromBody] CreateChatR
                 Temperature = s.Temperature,
                 EnableSearch = s.EnableSearch,
             }).ToArray(),
-            MessageCount = chat.Messages.Count,
+            LeafMessageId = chat.LeafMessageId != null ? idEncryption.EncryptMessageId(chat.LeafMessageId.Value) : null,
         });
     }
 
@@ -178,11 +178,13 @@ await db.Chats
     [HttpPut("{encryptedChatId}")]
     public async Task<IActionResult> UpdateChats(string encryptedChatId, [FromBody] UpdateChatsRequest request, CancellationToken cancellationToken)
     {
-        if (request.ModelId != null)
+        if (!ModelState.IsValid)
         {
-            return BadRequest("ModelId is not allowed to be updated anymore, please use ChatSpan update API.");
+            return BadRequest(ModelState);
         }
 
+        DecryptedUpdateChatsRequest req = request.Decrypt(idEncryption);
+
         Chat? chat = await db.Chats
             .Where(x => x.Id == idEncryption.DecryptChatId(encryptedChatId) && x.UserId == currentUser.Id)
             .FirstOrDefaultAsync(cancellationToken);
@@ -191,7 +193,13 @@ public async Task<IActionResult> UpdateChats(string encryptedChatId, [FromBody]
             return NotFound();
         }
 
-        request.ApplyToChats(chat);
+        string? error = await req.Validate(db, chat.Id);
+        if (error != null)
+        {
+            return BadRequest(error);
+        }
+
+        req.ApplyToChats(chat);
         if (db.ChangeTracker.HasChanges())
         {
             await db.SaveChangesAsync(cancellationToken);

diff --git a/src/BE/DB/Chat.cs b/src/BE/DB/Chat.cs
@@ -7,6 +7,7 @@
 namespace Chats.BE.DB;
 
 [Table("Chat")]
+[Index("CreatedAt", Name = "IX_Chat_CreateAt")]
 [Index("UserId", Name = "IX_Chat_UserId")]
 public partial class Chat
 {
@@ -20,13 +21,19 @@ public partial class Chat
 
     public bool IsDeleted { get; set; }
 
+    public long? LeafMessageId { get; set; }
+
     public DateTime CreatedAt { get; set; }
 
     public int UserId { get; set; }
 
     [InverseProperty("Chat")]
     public virtual ICollection<ChatSpan> ChatSpans { get; set; } = new List<ChatSpan>();
 
+    [ForeignKey("LeafMessageId")]
+    [InverseProperty("Chats")]
+    public virtual Message? LeafMessage { get; set; }
+
     [InverseProperty("Chat")]
     public virtual ICollection<Message> Messages { get; set; } = new List<Message>();
 

diff --git a/src/BE/DB/ChatsDB.cs b/src/BE/DB/ChatsDB.cs
@@ -116,7 +116,7 @@ protected override void OnModelCreating(ModelBuilder modelBuilder)
 
         modelBuilder.Entity<Chat>(entity =>
         {
-            entity.HasKey(e => e.Id).HasName("PK_Conversation2");
+            entity.HasOne(d => d.LeafMessage).WithMany(p => p.Chats).HasConstraintName("FK_Chat_Message");
 
             entity.HasOne(d => d.User).WithMany(p => p.Chats)
                 .OnDelete(DeleteBehavior.ClientSetNull)

diff --git a/src/BE/DB/Message.cs b/src/BE/DB/Message.cs
@@ -34,6 +34,9 @@ public partial class Message
     [InverseProperty("Messages")]
     public virtual ChatRole ChatRole { get; set; } = null!;
 
+    [InverseProperty("LeafMessage")]
+    public virtual ICollection<Chat> Chats { get; set; } = new List<Chat>();
+
     [InverseProperty("Parent")]
     public virtual ICollection<Message> InverseParent { get; set; } = new List<Message>();