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

…hat-span

src/BE/Controllers/Chats/Chats/ChatController.cs

@@ -67,17 +67,20 @@ public async Task<IActionResult> StartConversationStreamed(
         }
 
         // ensure chat.ChatSpan contains all span ids that in request, otherwise return error
-        if (req.Spans.Any(x => !chat.ChatSpans.Any(y => y.SpanId == x.Id)))
+        if (req.MessageId == null)
         {
-            return BadRequest("Invalid span id");
+            if (req.Spans.Any(x => !chat.ChatSpans.Any(y => y.SpanId == x.Id)))
+            {
+                return BadRequest("Invalid span id");
+            }
         }
 
         // get span id -> model id mapping but request only contains span id, so we need to get model id from chat.ChatSpan
         Dictionary<byte, short> spanModelMapping = req.Spans.ToDictionary(x => x.Id, x => chat.ChatSpans.First(y => y.SpanId == x.Id).ModelId);
         Dictionary<short, UserModel> userModels = await userModelManager.GetUserModels(currentUser.Id, [.. spanModelMapping.Values], cancellationToken);
 
-        // ensure all model ids are valid
-        if (userModels.Count != spanModelMapping.Count)
+        // ensure spanModelMapping contains all userModels
+        if (spanModelMapping.Values.Any(x => !userModels.ContainsKey(x)))
         {
             return BadRequest("Invalid span model");
         }
@@ -310,6 +313,7 @@ private static async Task<ChatSpanResponse> ProcessChatSpan(
             [
                 MessageContent.FromText(icc.FullResponse.TextSegment),
             ],
+            SpanId = span.Id,
             CreatedAt = DateTime.UtcNow,
             ParentId = req.MessageId,
         };

src/BE/Controllers/Chats/UserChats/UserChatsController.cs

@@ -99,18 +99,19 @@ public async Task<ActionResult<ChatsResponse>> CreateChat([FromBody] CreateChatR
         };
 
         Chat? lastChat = await db.Chats
-            .Include(x => x.ChatSpans)
+            .Include(x => x.ChatSpans.OrderBy(x => x.SpanId))
             .Where(x => x.UserId == currentUser.Id && !x.IsDeleted && x.ChatSpans.Any())
             .OrderByDescending(x => x.CreatedAt)
             .FirstOrDefaultAsync(cancellationToken);
         if (lastChat != null && lastChat.ChatSpans.All(cs => validModels.ContainsKey(cs.ModelId)))
         {
-            chat.ChatSpans = lastChat.ChatSpans.Select(cs => new ChatSpan()
+            chat.ChatSpans = lastChat.ChatSpans.Select((cs, i) => new ChatSpan()
             {
                 ModelId = cs.ModelId,
                 Model = validModels[cs.ModelId].Model,
                 EnableSearch = cs.EnableSearch,
                 Temperature = cs.Temperature,
+                SpanId = (byte)i,
             }).ToList();
         }
         else

src/BE/DB/ChatSpan.cs

@@ -27,9 +27,6 @@ public partial class ChatSpan
     [InverseProperty("ChatSpans")]
     public virtual Chat Chat { get; set; } = null!;
 
-    [InverseProperty("ChatSpan")]
-    public virtual ICollection<Message> Messages { get; set; } = new List<Message>();
-
     [ForeignKey("ModelId")]
     [InverseProperty("ChatSpans")]
     public virtual Model Model { get; set; } = null!;

src/BE/DB/ChatsDB.cs

@@ -209,8 +209,6 @@ protected override void OnModelCreating(ModelBuilder modelBuilder)
             entity.HasOne(d => d.Parent).WithMany(p => p.InverseParent).HasConstraintName("FK_Message_ParentMessage");
 
             entity.HasOne(d => d.Usage).WithOne(p => p.Message).HasConstraintName("FK_Message_UserModelUsage");
-
-            entity.HasOne(d => d.ChatSpan).WithMany(p => p.Messages).HasConstraintName("FK_Message_ChatSpan");
         });
 
         modelBuilder.Entity<MessageContent>(entity =>

src/BE/DB/Message.cs

@@ -34,10 +34,6 @@ public partial class Message
     [InverseProperty("Messages")]
     public virtual ChatRole ChatRole { get; set; } = null!;
 
-    [ForeignKey("ChatId, SpanId")]
-    [InverseProperty("Messages")]
-    public virtual ChatSpan? ChatSpan { get; set; }
-
     [InverseProperty("Parent")]
     public virtual ICollection<Message> InverseParent { get; set; } = new List<Message>();
 

src/scripts/db-migration/2024/20241217-chat-span.sql

@@ -865,4 +865,5 @@ ALTER TABLE dbo.MessageContent SET (LOCK_ESCALATION = TABLE)
 GO
 COMMIT
 
-update [Message] set [SpanId] = null where ChatRoleId = 2;
+update [Message] set [SpanId] = null where ChatRoleId = 2;
+ALTER TABLE dbo.Message DROP CONSTRAINT FK_Message_ChatSpan;