From 483adc59fe66d3573eadfd4b85335d6cae626b07 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Wed, 18 Mar 2026 12:27:07 -0400 Subject: [PATCH] feat: replace InsertChatMessage with batch InsertChatMessages (#23220) Replaces the singular `InsertChatMessage` query with `InsertChatMessages` that uses PostgreSQL's `unnest()` for batch inserts. This reduces the number of database round-trips when inserting multiple messages in a single transaction. ## Changes - **SQL**: New `InsertChatMessages :many` query using `unnest()` arrays following the existing codebase pattern (e.g., `InsertWorkspaceAgentStats`). Preserves the CTE that updates `chats.last_model_config_id` using the last non-null model config from the batch. Uses `NULLIF` for UUID columns to handle NULL foreign keys. - **Go layers**: Updated `querier.go`, `dbauthz.go`, `dbmetrics/querymetrics.go`, `dbmock/dbmock.go`, and `queries.sql.go` to use the new batch signature (`[]ChatMessage` return type, array params). - **chatd.go**: All call sites converted to batch inserts: - **CreateChat**: System prompt + user message batched into one call - **persistStep**: Assistant message + tool messages batched into one call - **persistSummary**: Hidden summary + assistant + tool messages batched into one call - Single-message sites use the same API with single-element arrays - **Helper**: New `appendChatMessage` function simplifies building batch params at each call site. - **Tests**: All test files updated to use the new API. Builds on top of #23213. --- coderd/chatd/chatd.go | 467 ++++++++++------------ coderd/chatd/chatd_test.go | 230 ++++++----- coderd/chats_test.go | 218 ++++++---- coderd/database/dbauthz/dbauthz.go | 8 +- coderd/database/dbauthz/dbauthz_test.go | 10 +- coderd/database/dbmetrics/querymetrics.go | 8 +- coderd/database/dbmock/dbmock.go | 14 +- coderd/database/querier.go | 2 +- coderd/database/querier_test.go | 61 +-- coderd/database/queries.sql.go | 196 +++++---- coderd/database/queries/chats.sql | 62 ++- 11 files changed, 692 insertions(+), 584 deletions(-) diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index 30aa4d99ad..94650b51e6 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -459,40 +459,6 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C } systemPrompt := strings.TrimSpace(opts.SystemPrompt) - if systemPrompt != "" { - systemContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText(systemPrompt), - }) - if err != nil { - return xerrors.Errorf("marshal system prompt: %w", err) - } - _, err = tx.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: insertedChat.ID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{ - UUID: opts.ModelConfigID, - Valid: true, - }, - Role: database.ChatMessageRoleSystem, - ContentVersion: chatprompt.CurrentContentVersion, - Content: systemContent, - Visibility: database.ChatMessageVisibilityModel, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - TotalCostMicros: sql.NullInt64{}, - RuntimeMs: sql.NullInt64{}, - }) - if err != nil { - return xerrors.Errorf("insert system message: %w", err) - } - } - var workspaceAwareness string if opts.WorkspaceID.Valid { workspaceAwareness = "This chat is attached to a workspace. You can use workspace tools like execute, read_file, write_file, etc." @@ -505,60 +471,53 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C if err != nil { return xerrors.Errorf("marshal workspace awareness: %w", err) } - _, err = tx.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: insertedChat.ID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{ - UUID: opts.ModelConfigID, - Valid: true, - }, - Role: database.ChatMessageRoleSystem, - ContentVersion: chatprompt.CurrentContentVersion, - Content: workspaceAwarenessContent, - Visibility: database.ChatMessageVisibilityModel, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - TotalCostMicros: sql.NullInt64{}, - RuntimeMs: sql.NullInt64{}, - }) - if err != nil { - return xerrors.Errorf("insert workspace awareness message: %w", err) - } - userContent, err := chatprompt.MarshalParts(opts.InitialUserContent) if err != nil { return xerrors.Errorf("marshal initial user content: %w", err) } - _, err = insertChatMessageWithStore(ctx, tx, database.InsertChatMessageParams{ + + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. ChatID: insertedChat.ID, - ModelConfigID: uuid.NullUUID{ - UUID: opts.ModelConfigID, - Valid: true, - }, - Role: database.ChatMessageRoleUser, - ContentVersion: chatprompt.CurrentContentVersion, - Content: userContent, - CreatedBy: uuid.NullUUID{UUID: opts.OwnerID, Valid: opts.OwnerID != uuid.Nil}, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - TotalCostMicros: sql.NullInt64{}, - RuntimeMs: sql.NullInt64{}, - Compressed: sql.NullBool{}, + } + + if systemPrompt != "" { + systemContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(systemPrompt), + }) + if err != nil { + return xerrors.Errorf("marshal system prompt: %w", err) + } + appendChatMessage(&msgParams, chatMessage{ + Role: database.ChatMessageRoleSystem, + Content: systemContent, + Visibility: database.ChatMessageVisibilityModel, + ModelConfigID: opts.ModelConfigID, + CreatedBy: uuid.Nil, + ContentVersion: chatprompt.CurrentContentVersion, + }) + } + + appendChatMessage(&msgParams, chatMessage{ + Role: database.ChatMessageRoleSystem, + Content: workspaceAwarenessContent, + Visibility: database.ChatMessageVisibilityModel, + ModelConfigID: opts.ModelConfigID, + CreatedBy: uuid.Nil, + ContentVersion: chatprompt.CurrentContentVersion, }) + + appendChatMessage(&msgParams, chatMessage{ + Role: database.ChatMessageRoleUser, + Content: userContent, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: opts.ModelConfigID, + CreatedBy: opts.OwnerID, + ContentVersion: chatprompt.CurrentContentVersion, + }) + + _, err = tx.InsertChatMessages(ctx, msgParams) if err != nil { - return xerrors.Errorf("insert initial user message: %w", err) + return xerrors.Errorf("insert initial chat messages: %w", err) } chat, err = setChatPendingWithStore(ctx, tx, insertedChat.ID) @@ -1140,13 +1099,60 @@ func (p *Server) setChatWaiting(ctx context.Context, chatID uuid.UUID) (database func insertChatMessageWithStore( ctx context.Context, store database.Store, - params database.InsertChatMessageParams, -) (database.ChatMessage, error) { - message, err := store.InsertChatMessage(ctx, params) + params database.InsertChatMessagesParams, +) ([]database.ChatMessage, error) { + messages, err := store.InsertChatMessages(ctx, params) if err != nil { - return database.ChatMessage{}, xerrors.Errorf("insert chat message: %w", err) + return nil, xerrors.Errorf("insert chat message: %w", err) } - return message, nil + return messages, nil +} + +// chatMessage describes a single message to insert as part of a batch. +// For nullable UUID fields (ModelConfigID, CreatedBy), use uuid.Nil to +// represent NULL — the SQL uses NULLIF to convert zero UUIDs to NULL. +// For nullable int64 fields, use 0 to represent NULL — the SQL uses +// NULLIF to convert zeros to NULL. +type chatMessage struct { + Role database.ChatMessageRole + Content pqtype.NullRawMessage + Visibility database.ChatMessageVisibility + ModelConfigID uuid.UUID + CreatedBy uuid.UUID + ContentVersion int16 + Compressed bool + InputTokens int64 + OutputTokens int64 + TotalTokens int64 + ReasoningTokens int64 + CacheCreationTokens int64 + CacheReadTokens int64 + ContextLimit int64 + TotalCostMicros int64 + RuntimeMs int64 +} + +// appendChatMessage appends a single message to the batch insert params. +func appendChatMessage( + params *database.InsertChatMessagesParams, + msg chatMessage, +) { + params.CreatedBy = append(params.CreatedBy, msg.CreatedBy) + params.ModelConfigID = append(params.ModelConfigID, msg.ModelConfigID) + params.Role = append(params.Role, msg.Role) + params.Content = append(params.Content, string(msg.Content.RawMessage)) + params.ContentVersion = append(params.ContentVersion, msg.ContentVersion) + params.Visibility = append(params.Visibility, msg.Visibility) + params.InputTokens = append(params.InputTokens, msg.InputTokens) + params.OutputTokens = append(params.OutputTokens, msg.OutputTokens) + params.TotalTokens = append(params.TotalTokens, msg.TotalTokens) + params.ReasoningTokens = append(params.ReasoningTokens, msg.ReasoningTokens) + params.CacheCreationTokens = append(params.CacheCreationTokens, msg.CacheCreationTokens) + params.CacheReadTokens = append(params.CacheReadTokens, msg.CacheReadTokens) + params.ContextLimit = append(params.ContextLimit, msg.ContextLimit) + params.Compressed = append(params.Compressed, msg.Compressed) + params.TotalCostMicros = append(params.TotalCostMicros, msg.TotalCostMicros) + params.RuntimeMs = append(params.RuntimeMs, msg.RuntimeMs) } func insertUserMessageAndSetPending( @@ -1157,28 +1163,22 @@ func insertUserMessageAndSetPending( content pqtype.NullRawMessage, createdBy uuid.UUID, ) (database.ChatMessage, database.Chat, error) { - message, err := insertChatMessageWithStore(ctx, store, database.InsertChatMessageParams{ - ChatID: lockedChat.ID, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: database.ChatMessageRoleUser, - ContentVersion: chatprompt.CurrentContentVersion, - Content: content, - CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: createdBy != uuid.Nil}, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - TotalCostMicros: sql.NullInt64{}, - RuntimeMs: sql.NullInt64{}, - Compressed: sql.NullBool{}, + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + ChatID: lockedChat.ID, + } + appendChatMessage(&msgParams, chatMessage{ + Role: database.ChatMessageRoleUser, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: modelConfigID, + CreatedBy: createdBy, + ContentVersion: chatprompt.CurrentContentVersion, }) + messages, err := insertChatMessageWithStore(ctx, store, msgParams) if err != nil { return database.ChatMessage{}, database.Chat{}, err } + message := messages[0] if lockedChat.Status == database.ChatStatusPending { return message, lockedChat, nil @@ -2132,33 +2132,27 @@ func (p *Server) tryAutoPromoteQueuedMessage( return nil, nil, false, xerrors.Errorf("pop next queued message: %w", err) } - msg, err := insertChatMessageWithStore(ctx, tx, database.InsertChatMessageParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: true}, - Role: database.ChatMessageRoleUser, - ContentVersion: chatprompt.CurrentContentVersion, + msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + ChatID: chat.ID, + } + appendChatMessage(&msgParams, chatMessage{ + Role: database.ChatMessageRoleUser, Content: pqtype.NullRawMessage{ RawMessage: nextQueued.Content, Valid: len(nextQueued.Content) > 0, }, - CreatedBy: uuid.NullUUID{UUID: chat.OwnerID, Valid: chat.OwnerID != uuid.Nil}, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - TotalCostMicros: sql.NullInt64{}, - RuntimeMs: sql.NullInt64{}, - Compressed: sql.NullBool{}, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: chat.LastModelConfigID, + CreatedBy: chat.OwnerID, + ContentVersion: chatprompt.CurrentContentVersion, }) + msgs, err := insertChatMessageWithStore(ctx, tx, msgParams) if err != nil { logger.Error(ctx, "failed to promote queued message", slog.F("queued_message_id", nextQueued.ID), slog.Error(err)) return nil, nil, false, nil } + msg := msgs[0] remainingQueuedMessages, err := tx.GetChatQueuedMessages(ctx, chat.ID) if err != nil { @@ -2690,65 +2684,73 @@ func (p *Server) runChat( } } - if assistantContent.Valid { - assistantMessage, insertErr := tx.InsertChatMessage(persistCtx, database.InsertChatMessageParams{ - ChatID: chat.ID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: usageNullInt64(step.Usage.InputTokens, hasUsage), - OutputTokens: usageNullInt64(step.Usage.OutputTokens, hasUsage), - TotalTokens: usageNullInt64(step.Usage.TotalTokens, hasUsage), - ReasoningTokens: usageNullInt64( - step.Usage.ReasoningTokens, - hasUsage, - ), - CacheCreationTokens: usageNullInt64( - step.Usage.CacheCreationTokens, - hasUsage, - ), - CacheReadTokens: usageNullInt64(step.Usage.CacheReadTokens, hasUsage), - ContextLimit: step.ContextLimit, - Compressed: sql.NullBool{}, - TotalCostMicros: usageNullInt64Ptr(totalCostMicros), - RuntimeMs: sql.NullInt64{ - Int64: step.Runtime.Milliseconds(), - Valid: step.Runtime > 0, - }, - }) - if insertErr != nil { - return xerrors.Errorf("insert assistant message: %w", insertErr) - } - insertedMessages = append(insertedMessages, assistantMessage) + stepParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + ChatID: chat.ID, } - for i, resultContent := range toolResultContents { - toolMessage, insertErr := tx.InsertChatMessage(persistCtx, database.InsertChatMessageParams{ - ChatID: chat.ID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: database.ChatMessageRoleTool, - ContentVersion: chatprompt.CurrentContentVersion, - Content: resultContent, + var contextLimit int64 + if step.ContextLimit.Valid { + contextLimit = step.ContextLimit.Int64 + } + + var runtimeMs int64 + if step.Runtime > 0 { + runtimeMs = step.Runtime.Milliseconds() + } + + var totalCostVal int64 + if totalCostMicros != nil { + totalCostVal = *totalCostMicros + } + + var inputTokens, outputTokens, totalTokens int64 + var reasoningTokens, cacheCreationTokens, cacheReadTokens int64 + if hasUsage { + inputTokens = step.Usage.InputTokens + outputTokens = step.Usage.OutputTokens + totalTokens = step.Usage.TotalTokens + reasoningTokens = step.Usage.ReasoningTokens + cacheCreationTokens = step.Usage.CacheCreationTokens + cacheReadTokens = step.Usage.CacheReadTokens + } + + if assistantContent.Valid { + appendChatMessage(&stepParams, chatMessage{ + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - TotalCostMicros: sql.NullInt64{}, - RuntimeMs: sql.NullInt64{}, - Compressed: sql.NullBool{}, + ModelConfigID: modelConfig.ID, + CreatedBy: uuid.Nil, + ContentVersion: chatprompt.CurrentContentVersion, + InputTokens: inputTokens, + OutputTokens: outputTokens, + TotalTokens: totalTokens, + ReasoningTokens: reasoningTokens, + CacheCreationTokens: cacheCreationTokens, + CacheReadTokens: cacheReadTokens, + ContextLimit: contextLimit, + TotalCostMicros: totalCostVal, + RuntimeMs: runtimeMs, }) + } + + for _, resultContent := range toolResultContents { + appendChatMessage(&stepParams, chatMessage{ + Role: database.ChatMessageRoleTool, + Content: resultContent, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: modelConfig.ID, + CreatedBy: uuid.Nil, + ContentVersion: chatprompt.CurrentContentVersion, + }) + } + + if len(stepParams.Role) > 0 { + inserted, insertErr := tx.InsertChatMessages(persistCtx, stepParams) if insertErr != nil { - return xerrors.Errorf("insert tool result %d: %w", i, insertErr) + return xerrors.Errorf("insert step messages: %w", insertErr) } - insertedMessages = append(insertedMessages, toolMessage) + insertedMessages = append(insertedMessages, inserted...) } return nil @@ -3110,82 +3112,51 @@ func (p *Server) persistChatContextSummary( var insertedMessages []database.ChatMessage txErr := p.db.InTx(func(tx database.Store) error { - _, txErr := tx.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chatID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: database.ChatMessageRoleUser, - ContentVersion: chatprompt.CurrentContentVersion, - Content: systemContent, - Visibility: database.ChatMessageVisibilityModel, - Compressed: sql.NullBool{Bool: true, Valid: true}, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - TotalCostMicros: sql.NullInt64{}, - RuntimeMs: sql.NullInt64{}, - }) - if txErr != nil { - return xerrors.Errorf("insert hidden summary message: %w", txErr) + summaryParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + ChatID: chatID, } - assistantMessage, txErr := tx.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chatID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: database.ChatMessageRoleAssistant, + // Hidden summary user message (not published to subscribers). + appendChatMessage(&summaryParams, chatMessage{ + Role: database.ChatMessageRoleUser, + Content: systemContent, + Visibility: database.ChatMessageVisibilityModel, + ModelConfigID: modelConfigID, + CreatedBy: uuid.Nil, ContentVersion: chatprompt.CurrentContentVersion, + Compressed: true, + }) + + // Assistant tool-call message. + appendChatMessage(&summaryParams, chatMessage{ + Role: database.ChatMessageRoleAssistant, Content: assistantContent, Visibility: database.ChatMessageVisibilityUser, - Compressed: sql.NullBool{ - Bool: true, - Valid: true, - }, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - TotalCostMicros: sql.NullInt64{}, - RuntimeMs: sql.NullInt64{}, - }) - if txErr != nil { - return xerrors.Errorf("insert summary tool call message: %w", txErr) - } - insertedMessages = append(insertedMessages, assistantMessage) - - toolMessage, txErr := tx.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chatID, - CreatedBy: uuid.NullUUID{}, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: database.ChatMessageRoleTool, + ModelConfigID: modelConfigID, + CreatedBy: uuid.Nil, ContentVersion: chatprompt.CurrentContentVersion, + Compressed: true, + }) + + // Tool result message. + appendChatMessage(&summaryParams, chatMessage{ + Role: database.ChatMessageRoleTool, Content: toolResult, Visibility: database.ChatMessageVisibilityBoth, - Compressed: sql.NullBool{ - Bool: true, - Valid: true, - }, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - TotalCostMicros: sql.NullInt64{}, - RuntimeMs: sql.NullInt64{}, + ModelConfigID: modelConfigID, + CreatedBy: uuid.Nil, + ContentVersion: chatprompt.CurrentContentVersion, + Compressed: true, }) + + allInserted, txErr := tx.InsertChatMessages(ctx, summaryParams) if txErr != nil { - return xerrors.Errorf("insert summary tool result message: %w", txErr) + return xerrors.Errorf("insert summary messages: %w", txErr) } - insertedMessages = append(insertedMessages, toolMessage) + // Skip the first message (hidden summary user msg) when + // publishing — only the assistant and tool messages are + // visible to subscribers. + insertedMessages = allInserted[1:] return nil }, nil) @@ -3297,24 +3268,6 @@ func int64Ptr(value int64) *int64 { return &value } -//nolint:revive // Boolean controls SQL NULL validity. -func usageNullInt64(value int64, valid bool) sql.NullInt64 { - if !valid { - return sql.NullInt64{} - } - return sql.NullInt64{ - Int64: value, - Valid: valid, - } -} - -func usageNullInt64Ptr(v *int64) sql.NullInt64 { - if v == nil { - return sql.NullInt64{} - } - return sql.NullInt64{Int64: *v, Valid: true} -} - func refreshChatWorkspaceSnapshot( ctx context.Context, chat database.Chat, diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index db5f362e05..2f562062ba 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -710,23 +710,24 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) { }) require.NoError(t, err) - _, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{ + _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: existingChat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, - RuntimeMs: sql.NullInt64{}, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{model.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Content: []string{string(assistantContent.RawMessage)}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{100}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) @@ -809,23 +810,24 @@ func TestPromoteQueuedAllowsAlreadyQueuedMessageWhenUsageLimitReached(t *testing }) require.NoError(t, err) - _, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{ + _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, - RuntimeMs: sql.NullInt64{}, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{model.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Content: []string{string(assistantContent.RawMessage)}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{100}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) @@ -998,23 +1000,24 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { }) require.NoError(t, err) - _, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{ + _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: spendChat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, - RuntimeMs: sql.NullInt64{}, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{model.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Content: []string{string(assistantContent.RawMessage)}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{100}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) @@ -1094,23 +1097,24 @@ func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) { }) require.NoError(t, err) - _, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{ + _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, - RuntimeMs: sql.NullInt64{}, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{model.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Content: []string{string(assistantContent.RawMessage)}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{100}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) @@ -1185,24 +1189,27 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) { }) require.NoError(t, err) - assistantMessage, err := db.InsertChatMessage(ctx, database.InsertChatMessageParams{ + assistantMessages, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - RuntimeMs: sql.NullInt64{}, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{model.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Content: []string{string(assistantContent.RawMessage)}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{0}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) + assistantMessage := assistantMessages[0] _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, @@ -1534,46 +1541,51 @@ func TestSubscribeAfterMessageID(t *testing.T) { }) require.NoError(t, err) - msg2, err := db.InsertChatMessage(ctx, database.InsertChatMessageParams{ + msg2Results, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: secondContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - RuntimeMs: sql.NullInt64{}, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{model.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Content: []string{string(secondContent.RawMessage)}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{0}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) + msg2 := msg2Results[0] thirdContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("third"), }) require.NoError(t, err) - _, err = db.InsertChatMessage(ctx, database.InsertChatMessageParams{ + _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleUser, - ContentVersion: chatprompt.CurrentContentVersion, - Content: thirdContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - RuntimeMs: sql.NullInt64{}, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{model.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Content: []string{string(thirdContent.RawMessage)}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{0}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) diff --git a/coderd/chats_test.go b/coderd/chats_test.go index 7e7a993779..221a2ed03e 100644 --- a/coderd/chats_test.go +++ b/coderd/chats_test.go @@ -127,22 +127,24 @@ func insertAssistantCostMessage( }) require.NoError(t, err) - _, err = db.InsertChatMessage(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessageParams{ + _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ ChatID: chatID, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{}, - OutputTokens: sql.NullInt64{}, - TotalTokens: sql.NullInt64{}, - ReasoningTokens: sql.NullInt64{}, - CacheCreationTokens: sql.NullInt64{}, - CacheReadTokens: sql.NullInt64{}, - ContextLimit: sql.NullInt64{}, - Compressed: sql.NullBool{}, - TotalCostMicros: sql.NullInt64{Int64: totalCostMicros, Valid: true}, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{modelConfigID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Content: []string{string(assistantContent.RawMessage)}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{totalCostMicros}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) } @@ -4060,16 +4062,27 @@ func seedChatCostFixture(t *testing.T) chatCostTestFixture { var earliestCreatedAt time.Time var latestCreatedAt time.Time for i := 0; i < 2; i++ { - message, err := db.InsertChatMessage(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessageParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: "assistant", - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{Int64: 100, Valid: true}, - OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, - TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true}, + results, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{modelConfig.ID}, + Role: []database.ChatMessageRole{"assistant"}, + Content: []string{"null"}, + ContentVersion: []int16{0}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{100}, + OutputTokens: []int64{50}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{500}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) + message := results[0] if i == 0 || message.CreatedAt.Before(earliestCreatedAt) { earliestCreatedAt = message.CreatedAt } @@ -4162,16 +4175,27 @@ func TestChatCostSummary_AdminDrilldown(t *testing.T) { }) require.NoError(t, err) - message, err := db.InsertChatMessage(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessageParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: "assistant", - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{Int64: 200, Valid: true}, - OutputTokens: sql.NullInt64{Int64: 100, Valid: true}, - TotalCostMicros: sql.NullInt64{Int64: 750, Valid: true}, + results, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{modelConfig.ID}, + Role: []database.ChatMessageRole{"assistant"}, + Content: []string{"null"}, + ContentVersion: []int16{0}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{200}, + OutputTokens: []int64{100}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{750}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) + message := results[0] options := codersdk.ChatCostSummaryOptions{ // Pad the DB-assigned timestamp so the query window cannot race it. StartDate: message.CreatedAt.Add(-time.Minute), @@ -4217,14 +4241,24 @@ func TestChatCostUsers(t *testing.T) { Title: "admin chat", }) require.NoError(t, err) - _, err = db.InsertChatMessage(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessageParams{ - ChatID: adminChat.ID, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: "assistant", - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{Int64: 100, Valid: true}, - OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, - TotalCostMicros: sql.NullInt64{Int64: 300, Valid: true}, + _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessagesParams{ + ChatID: adminChat.ID, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{modelConfig.ID}, + Role: []database.ChatMessageRole{"assistant"}, + Content: []string{"null"}, + ContentVersion: []int16{0}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{100}, + OutputTokens: []int64{50}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{300}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) @@ -4234,14 +4268,24 @@ func TestChatCostUsers(t *testing.T) { Title: "member chat", }) require.NoError(t, err) - _, err = db.InsertChatMessage(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessageParams{ - ChatID: memberChat.ID, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: "assistant", - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{Int64: 200, Valid: true}, - OutputTokens: sql.NullInt64{Int64: 100, Valid: true}, - TotalCostMicros: sql.NullInt64{Int64: 800, Valid: true}, + _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessagesParams{ + ChatID: memberChat.ID, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{modelConfig.ID}, + Role: []database.ChatMessageRole{"assistant"}, + Content: []string{"null"}, + ContentVersion: []int16{0}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{200}, + OutputTokens: []int64{100}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{800}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) @@ -4308,14 +4352,24 @@ func TestChatCostSummary_DateRange(t *testing.T) { }) require.NoError(t, err) - _, err = db.InsertChatMessage(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessageParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: "assistant", - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{Int64: 100, Valid: true}, - OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, - TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true}, + _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(seedCtx), database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{modelConfig.ID}, + Role: []database.ChatMessageRole{"assistant"}, + Content: []string{"null"}, + ContentVersion: []int16{0}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{100}, + OutputTokens: []int64{50}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{500}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) @@ -4363,27 +4417,49 @@ func TestChatCostSummary_UnpricedMessages(t *testing.T) { }) require.NoError(t, err) - pricedMessage, err := db.InsertChatMessage(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessageParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: "assistant", - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{Int64: 100, Valid: true}, - OutputTokens: sql.NullInt64{Int64: 50, Valid: true}, - TotalCostMicros: sql.NullInt64{Int64: 500, Valid: true}, + pricedResults, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{modelConfig.ID}, + Role: []database.ChatMessageRole{"assistant"}, + Content: []string{"null"}, + ContentVersion: []int16{0}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{100}, + OutputTokens: []int64{50}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{500}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) + pricedMessage := pricedResults[0] - unpricedMessage, err := db.InsertChatMessage(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessageParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, - Role: "assistant", - Visibility: database.ChatMessageVisibilityBoth, - InputTokens: sql.NullInt64{Int64: 200, Valid: true}, - OutputTokens: sql.NullInt64{Int64: 75, Valid: true}, - TotalCostMicros: sql.NullInt64{}, + unpricedResults, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{modelConfig.ID}, + Role: []database.ChatMessageRole{"assistant"}, + Content: []string{"null"}, + ContentVersion: []int16{0}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{200}, + OutputTokens: []int64{75}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{0}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) + unpricedMessage := unpricedResults[0] earliestCreatedAt := pricedMessage.CreatedAt latestCreatedAt := pricedMessage.CreatedAt diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index de5f2a6d27..9f4976efa8 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -4603,16 +4603,16 @@ func (q *querier) InsertChatFile(ctx context.Context, arg database.InsertChatFil return insert(q.log, q.auth, rbac.ResourceChat.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), q.db.InsertChatFile)(ctx, arg) } -func (q *querier) InsertChatMessage(ctx context.Context, arg database.InsertChatMessageParams) (database.ChatMessage, error) { +func (q *querier) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { // Authorize create on the parent chat (using update permission). chat, err := q.db.GetChatByID(ctx, arg.ChatID) if err != nil { - return database.ChatMessage{}, err + return nil, err } if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { - return database.ChatMessage{}, err + return nil, err } - return q.db.InsertChatMessage(ctx, arg) + return q.db.InsertChatMessages(ctx, arg) } func (q *querier) InsertChatModelConfig(ctx context.Context, arg database.InsertChatModelConfigParams) (database.ChatModelConfig, error) { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 0294a0dd3a..d11b349a09 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -675,13 +675,13 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().InsertChatFile(gomock.Any(), arg).Return(file, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), policy.ActionCreate).Returns(file) })) - s.Run("InsertChatMessage", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + s.Run("InsertChatMessages", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) - arg := testutil.Fake(s.T(), faker, database.InsertChatMessageParams{ChatID: chat.ID}) - msg := testutil.Fake(s.T(), faker, database.ChatMessage{ChatID: chat.ID}) + arg := testutil.Fake(s.T(), faker, database.InsertChatMessagesParams{ChatID: chat.ID}) + msgs := []database.ChatMessage{testutil.Fake(s.T(), faker, database.ChatMessage{ChatID: chat.ID})} dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() - dbm.EXPECT().InsertChatMessage(gomock.Any(), arg).Return(msg, nil).AnyTimes() - check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(msg) + dbm.EXPECT().InsertChatMessages(gomock.Any(), arg).Return(msgs, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(msgs) })) s.Run("InsertChatQueuedMessage", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 288f7c07e1..fceae90d74 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3064,11 +3064,11 @@ func (m queryMetricsStore) InsertChatFile(ctx context.Context, arg database.Inse return r0, r1 } -func (m queryMetricsStore) InsertChatMessage(ctx context.Context, arg database.InsertChatMessageParams) (database.ChatMessage, error) { +func (m queryMetricsStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { start := time.Now() - r0, r1 := m.s.InsertChatMessage(ctx, arg) - m.queryLatencies.WithLabelValues("InsertChatMessage").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "InsertChatMessage").Inc() + r0, r1 := m.s.InsertChatMessages(ctx, arg) + m.queryLatencies.WithLabelValues("InsertChatMessages").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "InsertChatMessages").Inc() return r0, r1 } diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 7c60866a2c..3aa4d683e5 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -5735,19 +5735,19 @@ func (mr *MockStoreMockRecorder) InsertChatFile(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatFile", reflect.TypeOf((*MockStore)(nil).InsertChatFile), ctx, arg) } -// InsertChatMessage mocks base method. -func (m *MockStore) InsertChatMessage(ctx context.Context, arg database.InsertChatMessageParams) (database.ChatMessage, error) { +// InsertChatMessages mocks base method. +func (m *MockStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "InsertChatMessage", ctx, arg) - ret0, _ := ret[0].(database.ChatMessage) + ret := m.ctrl.Call(m, "InsertChatMessages", ctx, arg) + ret0, _ := ret[0].([]database.ChatMessage) ret1, _ := ret[1].(error) return ret0, ret1 } -// InsertChatMessage indicates an expected call of InsertChatMessage. -func (mr *MockStoreMockRecorder) InsertChatMessage(ctx, arg any) *gomock.Call { +// InsertChatMessages indicates an expected call of InsertChatMessages. +func (mr *MockStoreMockRecorder) InsertChatMessages(ctx, arg any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatMessage", reflect.TypeOf((*MockStore)(nil).InsertChatMessage), ctx, arg) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatMessages", reflect.TypeOf((*MockStore)(nil).InsertChatMessages), ctx, arg) } // InsertChatModelConfig mocks base method. diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 41bdf082c7..cc9885efa0 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -637,7 +637,7 @@ type sqlcQuerier interface { InsertAuditLog(ctx context.Context, arg InsertAuditLogParams) (AuditLog, error) InsertChat(ctx context.Context, arg InsertChatParams) (Chat, error) InsertChatFile(ctx context.Context, arg InsertChatFileParams) (InsertChatFileRow, error) - InsertChatMessage(ctx context.Context, arg InsertChatMessageParams) (ChatMessage, error) + InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) InsertChatModelConfig(ctx context.Context, arg InsertChatModelConfigParams) (ChatModelConfig, error) InsertChatProvider(ctx context.Context, arg InsertChatProviderParams) (ChatProvider, error) InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (ChatQueuedMessage, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 37d12150bd..5dd774e4cf 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -9404,7 +9404,7 @@ func TestInsertWorkspaceAgentDevcontainers(t *testing.T) { } } -func TestInsertChatMessage(t *testing.T) { +func TestInsertChatMessages(t *testing.T) { t.Parallel() insertModelConfig := func( @@ -9478,17 +9478,24 @@ func TestInsertChatMessage(t *testing.T) { insertMessage := func(t *testing.T, store database.Store, ctx context.Context, chatID, userID, modelConfigID uuid.UUID, content string) { t.Helper() - _, err := store.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chatID, - CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: database.ChatMessageRoleUser, - ContentVersion: chatprompt.CurrentContentVersion, - Visibility: database.ChatMessageVisibilityBoth, - Content: pqtype.NullRawMessage{ - RawMessage: json.RawMessage(fmt.Sprintf("%q", content)), - Valid: true, - }, + _, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + ChatID: chatID, + CreatedBy: []uuid.UUID{userID}, + ModelConfigID: []uuid.UUID{modelConfigID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + Content: []string{fmt.Sprintf("%q", content)}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{0}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) } @@ -9583,19 +9590,27 @@ func TestGetChatMessagesForPromptByChatID(t *testing.T) { content string, ) database.ChatMessage { t.Helper() - msg, err := db.InsertChatMessage(ctx, database.InsertChatMessageParams{ - ChatID: chatID, - Role: role, - ContentVersion: chatprompt.CurrentContentVersion, - Visibility: vis, - Compressed: sql.NullBool{Bool: compressed, Valid: true}, - Content: pqtype.NullRawMessage{ - RawMessage: json.RawMessage(`"` + content + `"`), - Valid: true, - }, + results, err := db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + ChatID: chatID, + CreatedBy: []uuid.UUID{uuid.Nil}, + ModelConfigID: []uuid.UUID{uuid.Nil}, + Role: []database.ChatMessageRole{role}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Visibility: []database.ChatMessageVisibility{vis}, + Compressed: []bool{compressed}, + Content: []string{`"` + content + `"`}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + TotalCostMicros: []int64{0}, + RuntimeMs: []int64{0}, }) require.NoError(t, err) - return msg + return results[0] } msgIDs := func(msgs []database.ChatMessage) []int64 { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 6dfe20d029..bf50d312b7 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -4841,16 +4841,34 @@ func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat return i, err } -const insertChatMessage = `-- name: InsertChatMessage :one +const insertChatMessages = `-- name: InsertChatMessages :many WITH updated_chat AS ( UPDATE chats SET - last_model_config_id = $3::uuid + last_model_config_id = ( + SELECT val + FROM unnest($3::uuid[]) + WITH ORDINALITY AS t(val, ord) + WHERE val != '00000000-0000-0000-0000-000000000000'::uuid + ORDER BY ord DESC + LIMIT 1 + ) WHERE id = $1::uuid - AND $3::uuid IS NOT NULL - AND chats.last_model_config_id IS DISTINCT FROM $3::uuid + AND EXISTS ( + SELECT 1 + FROM unnest($3::uuid[]) + WHERE unnest != '00000000-0000-0000-0000-000000000000'::uuid + ) + AND chats.last_model_config_id IS DISTINCT FROM ( + SELECT val + FROM unnest($3::uuid[]) + WITH ORDINALITY AS t(val, ord) + WHERE val != '00000000-0000-0000-0000-000000000000'::uuid + ORDER BY ord DESC + LIMIT 1 + ) ) INSERT INTO chat_messages ( chat_id, @@ -4870,92 +4888,108 @@ INSERT INTO chat_messages ( compressed, total_cost_micros, runtime_ms -) VALUES ( - $1::uuid, - $2::uuid, - $3::uuid, - $4::chat_message_role, - $5::jsonb, - $6::smallint, - $7::chat_message_visibility, - $8::bigint, - $9::bigint, - $10::bigint, - $11::bigint, - $12::bigint, - $13::bigint, - $14::bigint, - COALESCE($15::boolean, FALSE), - $16::bigint, - $17::bigint ) +SELECT + $1::uuid, + NULLIF(unnest($2::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), + NULLIF(unnest($3::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), + unnest($4::chat_message_role[]), + unnest($5::text[])::jsonb, + unnest($6::smallint[]), + unnest($7::chat_message_visibility[]), + NULLIF(unnest($8::bigint[]), 0), + NULLIF(unnest($9::bigint[]), 0), + NULLIF(unnest($10::bigint[]), 0), + NULLIF(unnest($11::bigint[]), 0), + NULLIF(unnest($12::bigint[]), 0), + NULLIF(unnest($13::bigint[]), 0), + NULLIF(unnest($14::bigint[]), 0), + unnest($15::boolean[]), + NULLIF(unnest($16::bigint[]), 0), + NULLIF(unnest($17::bigint[]), 0) RETURNING id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms ` -type InsertChatMessageParams struct { - ChatID uuid.UUID `db:"chat_id" json:"chat_id"` - CreatedBy uuid.NullUUID `db:"created_by" json:"created_by"` - ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` - Role ChatMessageRole `db:"role" json:"role"` - Content pqtype.NullRawMessage `db:"content" json:"content"` - ContentVersion int16 `db:"content_version" json:"content_version"` - Visibility ChatMessageVisibility `db:"visibility" json:"visibility"` - InputTokens sql.NullInt64 `db:"input_tokens" json:"input_tokens"` - OutputTokens sql.NullInt64 `db:"output_tokens" json:"output_tokens"` - TotalTokens sql.NullInt64 `db:"total_tokens" json:"total_tokens"` - ReasoningTokens sql.NullInt64 `db:"reasoning_tokens" json:"reasoning_tokens"` - CacheCreationTokens sql.NullInt64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` - CacheReadTokens sql.NullInt64 `db:"cache_read_tokens" json:"cache_read_tokens"` - ContextLimit sql.NullInt64 `db:"context_limit" json:"context_limit"` - Compressed sql.NullBool `db:"compressed" json:"compressed"` - TotalCostMicros sql.NullInt64 `db:"total_cost_micros" json:"total_cost_micros"` - RuntimeMs sql.NullInt64 `db:"runtime_ms" json:"runtime_ms"` +type InsertChatMessagesParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + CreatedBy []uuid.UUID `db:"created_by" json:"created_by"` + ModelConfigID []uuid.UUID `db:"model_config_id" json:"model_config_id"` + Role []ChatMessageRole `db:"role" json:"role"` + Content []string `db:"content" json:"content"` + ContentVersion []int16 `db:"content_version" json:"content_version"` + Visibility []ChatMessageVisibility `db:"visibility" json:"visibility"` + InputTokens []int64 `db:"input_tokens" json:"input_tokens"` + OutputTokens []int64 `db:"output_tokens" json:"output_tokens"` + TotalTokens []int64 `db:"total_tokens" json:"total_tokens"` + ReasoningTokens []int64 `db:"reasoning_tokens" json:"reasoning_tokens"` + CacheCreationTokens []int64 `db:"cache_creation_tokens" json:"cache_creation_tokens"` + CacheReadTokens []int64 `db:"cache_read_tokens" json:"cache_read_tokens"` + ContextLimit []int64 `db:"context_limit" json:"context_limit"` + Compressed []bool `db:"compressed" json:"compressed"` + TotalCostMicros []int64 `db:"total_cost_micros" json:"total_cost_micros"` + RuntimeMs []int64 `db:"runtime_ms" json:"runtime_ms"` } -func (q *sqlQuerier) InsertChatMessage(ctx context.Context, arg InsertChatMessageParams) (ChatMessage, error) { - row := q.db.QueryRowContext(ctx, insertChatMessage, +func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) { + rows, err := q.db.QueryContext(ctx, insertChatMessages, arg.ChatID, - arg.CreatedBy, - arg.ModelConfigID, - arg.Role, - arg.Content, - arg.ContentVersion, - arg.Visibility, - arg.InputTokens, - arg.OutputTokens, - arg.TotalTokens, - arg.ReasoningTokens, - arg.CacheCreationTokens, - arg.CacheReadTokens, - arg.ContextLimit, - arg.Compressed, - arg.TotalCostMicros, - arg.RuntimeMs, + pq.Array(arg.CreatedBy), + pq.Array(arg.ModelConfigID), + pq.Array(arg.Role), + pq.Array(arg.Content), + pq.Array(arg.ContentVersion), + pq.Array(arg.Visibility), + pq.Array(arg.InputTokens), + pq.Array(arg.OutputTokens), + pq.Array(arg.TotalTokens), + pq.Array(arg.ReasoningTokens), + pq.Array(arg.CacheCreationTokens), + pq.Array(arg.CacheReadTokens), + pq.Array(arg.ContextLimit), + pq.Array(arg.Compressed), + pq.Array(arg.TotalCostMicros), + pq.Array(arg.RuntimeMs), ) - var i ChatMessage - err := row.Scan( - &i.ID, - &i.ChatID, - &i.ModelConfigID, - &i.CreatedAt, - &i.Role, - &i.Content, - &i.Visibility, - &i.InputTokens, - &i.OutputTokens, - &i.TotalTokens, - &i.ReasoningTokens, - &i.CacheCreationTokens, - &i.CacheReadTokens, - &i.ContextLimit, - &i.Compressed, - &i.CreatedBy, - &i.ContentVersion, - &i.TotalCostMicros, - &i.RuntimeMs, - ) - return i, err + if err != nil { + return nil, err + } + defer rows.Close() + var items []ChatMessage + for rows.Next() { + var i ChatMessage + if err := rows.Scan( + &i.ID, + &i.ChatID, + &i.ModelConfigID, + &i.CreatedAt, + &i.Role, + &i.Content, + &i.Visibility, + &i.InputTokens, + &i.OutputTokens, + &i.TotalTokens, + &i.ReasoningTokens, + &i.CacheCreationTokens, + &i.CacheReadTokens, + &i.ContextLimit, + &i.Compressed, + &i.CreatedBy, + &i.ContentVersion, + &i.TotalCostMicros, + &i.RuntimeMs, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil } const insertChatQueuedMessage = `-- name: InsertChatQueuedMessage :one diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 679ea5b6f5..da581de9c0 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -178,16 +178,34 @@ INSERT INTO chats ( RETURNING *; --- name: InsertChatMessage :one +-- name: InsertChatMessages :many WITH updated_chat AS ( UPDATE chats SET - last_model_config_id = sqlc.narg('model_config_id')::uuid + last_model_config_id = ( + SELECT val + FROM unnest(@model_config_id::uuid[]) + WITH ORDINALITY AS t(val, ord) + WHERE val != '00000000-0000-0000-0000-000000000000'::uuid + ORDER BY ord DESC + LIMIT 1 + ) WHERE id = @chat_id::uuid - AND sqlc.narg('model_config_id')::uuid IS NOT NULL - AND chats.last_model_config_id IS DISTINCT FROM sqlc.narg('model_config_id')::uuid + AND EXISTS ( + SELECT 1 + FROM unnest(@model_config_id::uuid[]) + WHERE unnest != '00000000-0000-0000-0000-000000000000'::uuid + ) + AND chats.last_model_config_id IS DISTINCT FROM ( + SELECT val + FROM unnest(@model_config_id::uuid[]) + WITH ORDINALITY AS t(val, ord) + WHERE val != '00000000-0000-0000-0000-000000000000'::uuid + ORDER BY ord DESC + LIMIT 1 + ) ) INSERT INTO chat_messages ( chat_id, @@ -207,25 +225,25 @@ INSERT INTO chat_messages ( compressed, total_cost_micros, runtime_ms -) VALUES ( - @chat_id::uuid, - sqlc.narg('created_by')::uuid, - sqlc.narg('model_config_id')::uuid, - @role::chat_message_role, - sqlc.narg('content')::jsonb, - @content_version::smallint, - @visibility::chat_message_visibility, - sqlc.narg('input_tokens')::bigint, - sqlc.narg('output_tokens')::bigint, - sqlc.narg('total_tokens')::bigint, - sqlc.narg('reasoning_tokens')::bigint, - sqlc.narg('cache_creation_tokens')::bigint, - sqlc.narg('cache_read_tokens')::bigint, - sqlc.narg('context_limit')::bigint, - COALESCE(sqlc.narg('compressed')::boolean, FALSE), - sqlc.narg('total_cost_micros')::bigint, - sqlc.narg('runtime_ms')::bigint ) +SELECT + @chat_id::uuid, + NULLIF(unnest(@created_by::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), + NULLIF(unnest(@model_config_id::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), + unnest(@role::chat_message_role[]), + unnest(@content::text[])::jsonb, + unnest(@content_version::smallint[]), + unnest(@visibility::chat_message_visibility[]), + NULLIF(unnest(@input_tokens::bigint[]), 0), + NULLIF(unnest(@output_tokens::bigint[]), 0), + NULLIF(unnest(@total_tokens::bigint[]), 0), + NULLIF(unnest(@reasoning_tokens::bigint[]), 0), + NULLIF(unnest(@cache_creation_tokens::bigint[]), 0), + NULLIF(unnest(@cache_read_tokens::bigint[]), 0), + NULLIF(unnest(@context_limit::bigint[]), 0), + unnest(@compressed::boolean[]), + NULLIF(unnest(@total_cost_micros::bigint[]), 0), + NULLIF(unnest(@runtime_ms::bigint[]), 0) RETURNING *;