mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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.
This commit is contained in:
Kyle Carberry
+210
-257
@@ -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,
|
||||
|
||||
+121
-109
@@ -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)
|
||||
|
||||
|
||||
+147
-71
@@ -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
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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{})
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 {
|
||||
|
||||
+115
-81
@@ -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
|
||||
|
||||
@@ -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
|
||||
*;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user