mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add manual chat title regeneration (#23633)
## Summary
Adds a "Generate new title" action that lets users manually regenerate a
chat's title using richer conversation context than the automatic
first-message title path.
## Changes
### Backend
- **New endpoint:** `POST
/api/experimental/chats/{chatID}/title/regenerate` returns the updated
Chat with a regenerated title
- **Manual title algorithm:** Extracts useful user/assistant text turns
→ selects first user turn + last 3 turns → builds context with gap
markers → renders prompt with anti-recency guidance → calls lightweight
model → normalizes output
- **Helpers:** `extractManualTitleTurns`,
`selectManualTitleTurnIndexes`, `buildManualTitleContext`,
`renderManualTitlePrompt`, `generateManualTitle` — all private, with the
public `Server.RegenerateChatTitle` method
- **SDK:** `ExperimentalClient.RegenerateChatTitle(ctx, chatID) (Chat,
error)`
- Persists title via existing `UpdateChatByID` and broadcasts
`ChatEventKindTitleChange`
### Frontend
- API client method + React Query mutation with cache invalidation
- "Generate new title" menu item (with wand icon) in both TopBar and
Sidebar dropdown menus
- Loading/disabled state while regeneration is in-flight
- Error toast on failure
- Stories updated for both menus
### Tests
- `quickgen_test.go`: Table-driven tests for all 4 helper functions
(turn extraction, index selection, context building, prompt rendering)
- `exp_chats_test.go`: Handler tests (ChatNotFound,
NotFoundForDifferentUser, NoDaemon)
## Design notes
- The existing auto-title path (`maybeGenerateChatTitle`, `titleInput`)
is completely unchanged
- Manual regeneration uses richer context (first user turn + last 3
turns + gap markers) vs the auto path's single first message
- Endpoint is experimental and marked with `@x-apidocgen {"skip": true}`
This commit is contained in:
+499
-37
@@ -1333,6 +1333,477 @@ func (p *Server) InterruptChat(
|
||||
return updatedChat
|
||||
}
|
||||
|
||||
const manualTitleMessageWindowLimit = 50
|
||||
|
||||
var ErrManualTitleRegenerationInProgress = xerrors.New(
|
||||
"manual title regeneration already in progress",
|
||||
)
|
||||
|
||||
type manualTitleGenerationError struct {
|
||||
cause error
|
||||
modelConfig database.ChatModelConfig
|
||||
usage fantasy.Usage
|
||||
}
|
||||
|
||||
func (e *manualTitleGenerationError) Error() string {
|
||||
return e.cause.Error()
|
||||
}
|
||||
|
||||
func (e *manualTitleGenerationError) Unwrap() error {
|
||||
return e.cause
|
||||
}
|
||||
|
||||
var manualTitleLockWorkerID = uuid.MustParse(
|
||||
"00000000-0000-0000-0000-000000000001",
|
||||
)
|
||||
|
||||
const manualTitleLockStaleAfter = time.Minute
|
||||
|
||||
func isPendingOrRunningChatStatus(status database.ChatStatus) bool {
|
||||
switch status {
|
||||
case database.ChatStatusPending, database.ChatStatusRunning:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func isFreshManualTitleLock(chat database.Chat, now time.Time) bool {
|
||||
if !chat.WorkerID.Valid || chat.WorkerID.UUID != manualTitleLockWorkerID {
|
||||
return false
|
||||
}
|
||||
leaseAt := chat.HeartbeatAt
|
||||
if !leaseAt.Valid {
|
||||
leaseAt = chat.StartedAt
|
||||
}
|
||||
return leaseAt.Valid && leaseAt.Time.After(now.Add(-manualTitleLockStaleAfter))
|
||||
}
|
||||
|
||||
// updateChatStatusPreserveUpdatedAt applies internal lock transitions without
|
||||
// changing chat recency, because chat list ordering uses updated_at.
|
||||
func updateChatStatusPreserveUpdatedAt(
|
||||
ctx context.Context,
|
||||
store database.Store,
|
||||
chat database.Chat,
|
||||
workerID uuid.NullUUID,
|
||||
startedAt sql.NullTime,
|
||||
heartbeatAt sql.NullTime,
|
||||
) (database.Chat, error) {
|
||||
return store.UpdateChatStatusPreserveUpdatedAt(
|
||||
ctx,
|
||||
database.UpdateChatStatusPreserveUpdatedAtParams{
|
||||
ID: chat.ID,
|
||||
Status: chat.Status,
|
||||
WorkerID: workerID,
|
||||
StartedAt: startedAt,
|
||||
HeartbeatAt: heartbeatAt,
|
||||
LastError: chat.LastError,
|
||||
UpdatedAt: chat.UpdatedAt,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func (p *Server) acquireManualTitleLock(ctx context.Context, chatID uuid.UUID) error {
|
||||
now := time.Now()
|
||||
return p.db.InTx(func(tx database.Store) error {
|
||||
lockedChat, err := tx.GetChatByIDForUpdate(ctx, chatID)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("lock chat for manual title regeneration: %w", err)
|
||||
}
|
||||
if isPendingOrRunningChatStatus(lockedChat.Status) ||
|
||||
isFreshManualTitleLock(lockedChat, now) {
|
||||
return ErrManualTitleRegenerationInProgress
|
||||
}
|
||||
_, err = updateChatStatusPreserveUpdatedAt(
|
||||
ctx,
|
||||
tx,
|
||||
lockedChat,
|
||||
uuid.NullUUID{UUID: manualTitleLockWorkerID, Valid: true},
|
||||
sql.NullTime{Time: now, Valid: true},
|
||||
sql.NullTime{Time: now, Valid: true},
|
||||
)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("mark chat for manual title regeneration: %w", err)
|
||||
}
|
||||
return nil
|
||||
}, database.DefaultTXOptions().WithID("chat_title_regenerate_lock"))
|
||||
}
|
||||
|
||||
func (p *Server) releaseManualTitleLock(ctx context.Context, chatID uuid.UUID) {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
|
||||
defer cancel()
|
||||
|
||||
err := p.db.InTx(func(tx database.Store) error {
|
||||
lockedChat, err := tx.GetChatByIDForUpdate(cleanupCtx, chatID)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("lock chat to release manual title regeneration: %w", err)
|
||||
}
|
||||
if !lockedChat.WorkerID.Valid || lockedChat.WorkerID.UUID != manualTitleLockWorkerID {
|
||||
return nil
|
||||
}
|
||||
_, err = updateChatStatusPreserveUpdatedAt(
|
||||
cleanupCtx,
|
||||
tx,
|
||||
lockedChat,
|
||||
uuid.NullUUID{},
|
||||
sql.NullTime{},
|
||||
sql.NullTime{},
|
||||
)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("clear manual title regeneration marker: %w", err)
|
||||
}
|
||||
return nil
|
||||
}, database.DefaultTXOptions().WithID("chat_title_regenerate_unlock"))
|
||||
if err != nil {
|
||||
p.logger.Warn(cleanupCtx, "failed to release manual title regeneration marker",
|
||||
slog.F("chat_id", chatID),
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// RegenerateChatTitle regenerates a chat title from the chat's visible
|
||||
// messages, persists it when it changes, and broadcasts the update.
|
||||
func (p *Server) RegenerateChatTitle(
|
||||
ctx context.Context,
|
||||
chat database.Chat,
|
||||
) (database.Chat, error) {
|
||||
// Reuse chatd's scoped auth context for deployment-config lookups while
|
||||
// keeping chat ownership authorization at the HTTP layer.
|
||||
//nolint:gocritic // Non-admin users need chatd-scoped config reads here.
|
||||
chatdCtx := dbauthz.AsChatd(ctx)
|
||||
keys, err := p.resolveProviderAPIKeys(chatdCtx)
|
||||
if err != nil {
|
||||
return database.Chat{}, xerrors.Errorf("resolve chat providers: %w", err)
|
||||
}
|
||||
if err := p.acquireManualTitleLock(ctx, chat.ID); err != nil {
|
||||
return database.Chat{}, err
|
||||
}
|
||||
defer p.releaseManualTitleLock(chatdCtx, chat.ID)
|
||||
|
||||
updatedChat, err := p.regenerateChatTitleWithStore(
|
||||
chatdCtx,
|
||||
p.db,
|
||||
chat,
|
||||
keys,
|
||||
)
|
||||
if err != nil {
|
||||
var generationErr *manualTitleGenerationError
|
||||
if errors.As(err, &generationErr) {
|
||||
// Reuse chatd's scoped auth context for failure accounting while
|
||||
// detaching from request cancellation so usage is still recorded.
|
||||
//nolint:gocritic // Failure accounting still needs chatd-scoped config reads.
|
||||
recordCtx, recordCancel := context.WithTimeout(
|
||||
dbauthz.AsChatd(context.WithoutCancel(ctx)),
|
||||
5*time.Second,
|
||||
)
|
||||
defer recordCancel()
|
||||
if _, recordErr := recordManualTitleUsage(
|
||||
recordCtx,
|
||||
p.db,
|
||||
chat,
|
||||
generationErr.modelConfig,
|
||||
generationErr.usage,
|
||||
"",
|
||||
); recordErr != nil {
|
||||
return database.Chat{}, errors.Join(
|
||||
generationErr,
|
||||
xerrors.Errorf("record manual title usage: %w", recordErr),
|
||||
)
|
||||
}
|
||||
return database.Chat{}, generationErr
|
||||
}
|
||||
return database.Chat{}, err
|
||||
}
|
||||
return updatedChat, nil
|
||||
}
|
||||
|
||||
func (p *Server) regenerateChatTitleWithStore(
|
||||
ctx context.Context,
|
||||
store database.Store,
|
||||
chat database.Chat,
|
||||
keys chatprovider.ProviderAPIKeys,
|
||||
) (database.Chat, error) {
|
||||
if limitErr := p.checkUsageLimit(ctx, store, chat.OwnerID); limitErr != nil {
|
||||
return database.Chat{}, limitErr
|
||||
}
|
||||
|
||||
headMessages, err := store.GetChatMessagesByChatIDAscPaginated(
|
||||
ctx,
|
||||
database.GetChatMessagesByChatIDAscPaginatedParams{
|
||||
ChatID: chat.ID,
|
||||
AfterID: 0,
|
||||
LimitVal: manualTitleMessageWindowLimit,
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
return database.Chat{}, xerrors.Errorf("get head chat messages: %w", err)
|
||||
}
|
||||
tailMessages, err := store.GetChatMessagesByChatIDDescPaginated(
|
||||
ctx,
|
||||
database.GetChatMessagesByChatIDDescPaginatedParams{
|
||||
ChatID: chat.ID,
|
||||
BeforeID: 0,
|
||||
LimitVal: manualTitleMessageWindowLimit,
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
return database.Chat{}, xerrors.Errorf("get tail chat messages: %w", err)
|
||||
}
|
||||
messages := mergeManualTitleMessages(headMessages, tailMessages)
|
||||
if len(messages) == 0 {
|
||||
return chat, nil
|
||||
}
|
||||
|
||||
model, modelConfig, err := p.resolveManualTitleModel(ctx, store, chat, keys)
|
||||
if err != nil {
|
||||
return database.Chat{}, err
|
||||
}
|
||||
|
||||
title, usage, err := generateManualTitle(ctx, messages, model)
|
||||
if err != nil {
|
||||
wrappedErr := xerrors.Errorf("generate manual title: %w", err)
|
||||
if usage == (fantasy.Usage{}) {
|
||||
return database.Chat{}, wrappedErr
|
||||
}
|
||||
return database.Chat{}, &manualTitleGenerationError{
|
||||
cause: wrappedErr,
|
||||
modelConfig: modelConfig,
|
||||
usage: usage,
|
||||
}
|
||||
}
|
||||
|
||||
recordCtx, recordCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
|
||||
defer recordCancel()
|
||||
|
||||
updatedChat, recordErr := recordManualTitleUsage(
|
||||
recordCtx,
|
||||
store,
|
||||
chat,
|
||||
modelConfig,
|
||||
usage,
|
||||
title,
|
||||
)
|
||||
if recordErr != nil {
|
||||
if title != "" {
|
||||
return database.Chat{}, xerrors.Errorf("record manual title usage and update chat title: %w", recordErr)
|
||||
}
|
||||
return database.Chat{}, xerrors.Errorf("record manual title usage: %w", recordErr)
|
||||
}
|
||||
if updatedChat.Title == chat.Title {
|
||||
return updatedChat, nil
|
||||
}
|
||||
|
||||
p.publishChatPubsubEvent(updatedChat, coderdpubsub.ChatEventKindTitleChange, nil)
|
||||
return updatedChat, nil
|
||||
}
|
||||
|
||||
func (p *Server) resolveManualTitleModel(
|
||||
ctx context.Context,
|
||||
store database.Store,
|
||||
chat database.Chat,
|
||||
keys chatprovider.ProviderAPIKeys,
|
||||
) (fantasy.LanguageModel, database.ChatModelConfig, error) {
|
||||
configs, err := store.GetEnabledChatModelConfigs(ctx)
|
||||
if err != nil {
|
||||
p.logger.Debug(ctx, "failed to list manual title model configs",
|
||||
slog.F("chat_id", chat.ID),
|
||||
slog.Error(err),
|
||||
)
|
||||
return p.resolveFallbackManualTitleModel(ctx, chat, keys)
|
||||
}
|
||||
|
||||
config, ok := selectPreferredConfiguredShortTextModelConfig(configs)
|
||||
if !ok {
|
||||
return p.resolveFallbackManualTitleModel(ctx, chat, keys)
|
||||
}
|
||||
|
||||
model, err := chatprovider.ModelFromConfig(
|
||||
config.Provider,
|
||||
config.Model,
|
||||
keys,
|
||||
chatprovider.UserAgent(),
|
||||
chatprovider.CoderHeaders(chat),
|
||||
)
|
||||
if err != nil {
|
||||
p.logger.Debug(ctx, "manual title preferred model unavailable",
|
||||
slog.F("chat_id", chat.ID),
|
||||
slog.F("provider", config.Provider),
|
||||
slog.F("model", config.Model),
|
||||
slog.Error(err),
|
||||
)
|
||||
return p.resolveFallbackManualTitleModel(ctx, chat, keys)
|
||||
}
|
||||
|
||||
return model, config, nil
|
||||
}
|
||||
|
||||
func (p *Server) resolveFallbackManualTitleModel(
|
||||
ctx context.Context,
|
||||
chat database.Chat,
|
||||
keys chatprovider.ProviderAPIKeys,
|
||||
) (fantasy.LanguageModel, database.ChatModelConfig, error) {
|
||||
config, err := p.resolveModelConfig(ctx, chat)
|
||||
if err != nil {
|
||||
return nil, database.ChatModelConfig{}, xerrors.Errorf(
|
||||
"resolve fallback manual title model config: %w",
|
||||
err,
|
||||
)
|
||||
}
|
||||
model, err := chatprovider.ModelFromConfig(
|
||||
config.Provider,
|
||||
config.Model,
|
||||
keys,
|
||||
chatprovider.UserAgent(),
|
||||
chatprovider.CoderHeaders(chat),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, database.ChatModelConfig{}, xerrors.Errorf(
|
||||
"create fallback manual title model: %w",
|
||||
err,
|
||||
)
|
||||
}
|
||||
return model, config, nil
|
||||
}
|
||||
|
||||
func mergeManualTitleMessages(
|
||||
headMessages []database.ChatMessage,
|
||||
tailMessagesDesc []database.ChatMessage,
|
||||
) []database.ChatMessage {
|
||||
merged := make([]database.ChatMessage, 0, len(headMessages)+len(tailMessagesDesc))
|
||||
seen := make(map[int64]struct{}, len(headMessages)+len(tailMessagesDesc))
|
||||
appendUnique := func(message database.ChatMessage) {
|
||||
if _, ok := seen[message.ID]; ok {
|
||||
return
|
||||
}
|
||||
seen[message.ID] = struct{}{}
|
||||
merged = append(merged, message)
|
||||
}
|
||||
for _, message := range headMessages {
|
||||
appendUnique(message)
|
||||
}
|
||||
for i := len(tailMessagesDesc) - 1; i >= 0; i-- {
|
||||
appendUnique(tailMessagesDesc[i])
|
||||
}
|
||||
return merged
|
||||
}
|
||||
|
||||
func fantasyUsageToChatMessageUsage(usage fantasy.Usage) codersdk.ChatMessageUsage {
|
||||
var chatUsage codersdk.ChatMessageUsage
|
||||
if usage.InputTokens != 0 {
|
||||
chatUsage.InputTokens = ptr.Ref(usage.InputTokens)
|
||||
}
|
||||
if usage.OutputTokens != 0 {
|
||||
chatUsage.OutputTokens = ptr.Ref(usage.OutputTokens)
|
||||
}
|
||||
if usage.ReasoningTokens != 0 {
|
||||
chatUsage.ReasoningTokens = ptr.Ref(usage.ReasoningTokens)
|
||||
}
|
||||
if usage.CacheCreationTokens != 0 {
|
||||
chatUsage.CacheCreationTokens = ptr.Ref(usage.CacheCreationTokens)
|
||||
}
|
||||
if usage.CacheReadTokens != 0 {
|
||||
chatUsage.CacheReadTokens = ptr.Ref(usage.CacheReadTokens)
|
||||
}
|
||||
return chatUsage
|
||||
}
|
||||
|
||||
func recordManualTitleUsage(
|
||||
ctx context.Context,
|
||||
store database.Store,
|
||||
chat database.Chat,
|
||||
modelConfig database.ChatModelConfig,
|
||||
usage fantasy.Usage,
|
||||
newTitle string,
|
||||
) (database.Chat, error) {
|
||||
hasUsage := usage != (fantasy.Usage{})
|
||||
if !hasUsage && newTitle == "" {
|
||||
return chat, nil
|
||||
}
|
||||
|
||||
var totalCostMicros *int64
|
||||
if hasUsage {
|
||||
callConfig := codersdk.ChatModelCallConfig{}
|
||||
if len(modelConfig.Options) > 0 {
|
||||
if err := json.Unmarshal(modelConfig.Options, &callConfig); err != nil {
|
||||
return database.Chat{}, xerrors.Errorf("parse model call config: %w", err)
|
||||
}
|
||||
}
|
||||
totalCostMicros = chatcost.CalculateTotalCostMicros(
|
||||
fantasyUsageToChatMessageUsage(usage),
|
||||
callConfig.Cost,
|
||||
)
|
||||
}
|
||||
|
||||
// Use a valid empty JSON array for the content column.
|
||||
// MarshalParts returns a null NullRawMessage for empty
|
||||
// slices, which becomes an empty string that PostgreSQL
|
||||
// rejects as invalid JSON.
|
||||
content := "[]"
|
||||
|
||||
updatedChat := chat
|
||||
err := store.InTx(func(tx database.Store) error {
|
||||
lockedChat, err := tx.GetChatByIDForUpdate(ctx, chat.ID)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("lock chat for manual title usage: %w", err)
|
||||
}
|
||||
updatedChat = lockedChat
|
||||
if hasUsage {
|
||||
messages, err := tx.InsertChatMessages(ctx, database.InsertChatMessagesParams{
|
||||
ChatID: chat.ID,
|
||||
CreatedBy: []uuid.UUID{chat.OwnerID},
|
||||
ModelConfigID: []uuid.UUID{modelConfig.ID},
|
||||
Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant},
|
||||
Content: []string{content},
|
||||
ContentVersion: []int16{chatprompt.CurrentContentVersion},
|
||||
Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityModel},
|
||||
InputTokens: []int64{usage.InputTokens},
|
||||
OutputTokens: []int64{usage.OutputTokens},
|
||||
TotalTokens: []int64{usage.TotalTokens},
|
||||
ReasoningTokens: []int64{usage.ReasoningTokens},
|
||||
CacheCreationTokens: []int64{usage.CacheCreationTokens},
|
||||
CacheReadTokens: []int64{usage.CacheReadTokens},
|
||||
ContextLimit: []int64{modelConfig.ContextLimit},
|
||||
Compressed: []bool{false},
|
||||
TotalCostMicros: []int64{ptr.NilToDefault(totalCostMicros, 0)},
|
||||
RuntimeMs: []int64{0},
|
||||
ProviderResponseID: []string{""},
|
||||
})
|
||||
if err != nil {
|
||||
return xerrors.Errorf("insert manual title usage message: %w", err)
|
||||
}
|
||||
if len(messages) != 1 {
|
||||
return xerrors.Errorf("expected 1 manual title usage message, got %d", len(messages))
|
||||
}
|
||||
if err := tx.SoftDeleteChatMessageByID(ctx, messages[0].ID); err != nil {
|
||||
return xerrors.Errorf("soft delete manual title usage message: %w", err)
|
||||
}
|
||||
if lockedChat.LastModelConfigID != modelConfig.ID {
|
||||
if _, err := tx.UpdateChatLastModelConfigByID(ctx, database.UpdateChatLastModelConfigByIDParams{
|
||||
ID: chat.ID,
|
||||
LastModelConfigID: lockedChat.LastModelConfigID,
|
||||
}); err != nil {
|
||||
return xerrors.Errorf("restore chat model config after manual title usage: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
if newTitle != "" && lockedChat.Title == chat.Title && newTitle != lockedChat.Title {
|
||||
updatedChat, err = tx.UpdateChatByID(ctx, database.UpdateChatByIDParams{
|
||||
ID: chat.ID,
|
||||
Title: newTitle,
|
||||
})
|
||||
if err != nil {
|
||||
return xerrors.Errorf("update chat title: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}, nil)
|
||||
if err != nil {
|
||||
return database.Chat{}, err
|
||||
}
|
||||
return updatedChat, nil
|
||||
}
|
||||
|
||||
// RefreshStatus loads the latest chat status and publishes it to stream subscribers.
|
||||
func (p *Server) RefreshStatus(ctx context.Context, chatID uuid.UUID) error {
|
||||
if chatID == uuid.Nil {
|
||||
@@ -3475,24 +3946,7 @@ func (p *Server) runChat(
|
||||
}
|
||||
|
||||
hasUsage := step.Usage != (fantasy.Usage{})
|
||||
var usageForCost codersdk.ChatMessageUsage
|
||||
if hasUsage {
|
||||
if step.Usage.InputTokens != 0 {
|
||||
usageForCost.InputTokens = ptr.Ref(step.Usage.InputTokens)
|
||||
}
|
||||
if step.Usage.OutputTokens != 0 {
|
||||
usageForCost.OutputTokens = ptr.Ref(step.Usage.OutputTokens)
|
||||
}
|
||||
if step.Usage.ReasoningTokens != 0 {
|
||||
usageForCost.ReasoningTokens = ptr.Ref(step.Usage.ReasoningTokens)
|
||||
}
|
||||
if step.Usage.CacheCreationTokens != 0 {
|
||||
usageForCost.CacheCreationTokens = ptr.Ref(step.Usage.CacheCreationTokens)
|
||||
}
|
||||
if step.Usage.CacheReadTokens != 0 {
|
||||
usageForCost.CacheReadTokens = ptr.Ref(step.Usage.CacheReadTokens)
|
||||
}
|
||||
}
|
||||
usageForCost := fantasyUsageToChatMessageUsage(step.Usage)
|
||||
totalCostMicros := chatcost.CalculateTotalCostMicros(usageForCost, callConfig.Cost)
|
||||
|
||||
var insertedMessages []database.ChatMessage
|
||||
@@ -4078,10 +4532,8 @@ func (p *Server) resolveChatModel(
|
||||
ctx context.Context,
|
||||
chat database.Chat,
|
||||
) (fantasy.LanguageModel, database.ChatModelConfig, chatprovider.ProviderAPIKeys, error) {
|
||||
var (
|
||||
dbConfig database.ChatModelConfig
|
||||
providers []database.ChatProvider
|
||||
)
|
||||
var dbConfig database.ChatModelConfig
|
||||
var keys chatprovider.ProviderAPIKeys
|
||||
|
||||
var g errgroup.Group
|
||||
g.Go(func() error {
|
||||
@@ -4094,28 +4546,15 @@ func (p *Server) resolveChatModel(
|
||||
})
|
||||
g.Go(func() error {
|
||||
var err error
|
||||
providers, err = p.configCache.EnabledProviders(ctx)
|
||||
keys, err = p.resolveProviderAPIKeys(ctx)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("get enabled chat providers: %w", err)
|
||||
return xerrors.Errorf("resolve provider API keys: %w", err)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err := g.Wait(); err != nil {
|
||||
return nil, database.ChatModelConfig{}, chatprovider.ProviderAPIKeys{}, err
|
||||
}
|
||||
dbProviders := make(
|
||||
[]chatprovider.ConfiguredProvider, 0, len(providers),
|
||||
)
|
||||
for _, provider := range providers {
|
||||
dbProviders = append(dbProviders, chatprovider.ConfiguredProvider{
|
||||
Provider: provider.Provider,
|
||||
APIKey: provider.APIKey,
|
||||
BaseURL: provider.BaseUrl,
|
||||
})
|
||||
}
|
||||
keys := chatprovider.MergeProviderAPIKeys(
|
||||
p.providerAPIKeys, dbProviders,
|
||||
)
|
||||
|
||||
model, err := chatprovider.ModelFromConfig(
|
||||
dbConfig.Provider, dbConfig.Model, keys, chatprovider.UserAgent(),
|
||||
@@ -4129,6 +4568,29 @@ func (p *Server) resolveChatModel(
|
||||
return model, dbConfig, keys, nil
|
||||
}
|
||||
|
||||
func (p *Server) resolveProviderAPIKeys(
|
||||
ctx context.Context,
|
||||
) (chatprovider.ProviderAPIKeys, error) {
|
||||
providers, err := p.configCache.EnabledProviders(ctx)
|
||||
if err != nil {
|
||||
return chatprovider.ProviderAPIKeys{}, xerrors.Errorf(
|
||||
"get enabled chat providers: %w",
|
||||
err,
|
||||
)
|
||||
}
|
||||
dbProviders := make(
|
||||
[]chatprovider.ConfiguredProvider, 0, len(providers),
|
||||
)
|
||||
for _, provider := range providers {
|
||||
dbProviders = append(dbProviders, chatprovider.ConfiguredProvider{
|
||||
Provider: provider.Provider,
|
||||
APIKey: provider.APIKey,
|
||||
BaseURL: provider.BaseUrl,
|
||||
})
|
||||
}
|
||||
return chatprovider.MergeProviderAPIKeys(p.providerAPIKeys, dbProviders), nil
|
||||
}
|
||||
|
||||
// resolveModelConfig looks up the chat's model config by its
|
||||
// LastModelConfigID. If the referenced config no longer exists
|
||||
// (e.g. it was deleted), it falls back to the default model
|
||||
|
||||
@@ -23,6 +23,7 @@ import (
|
||||
dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub"
|
||||
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chaterror"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chattest"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock"
|
||||
@@ -30,6 +31,183 @@ import (
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
func TestRegenerateChatTitle_PersistsAndBroadcasts(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
ctrl := gomock.NewController(t)
|
||||
db := dbmock.NewMockStore(ctrl)
|
||||
lockTx := dbmock.NewMockStore(ctrl)
|
||||
usageTx := dbmock.NewMockStore(ctrl)
|
||||
unlockTx := dbmock.NewMockStore(ctrl)
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
pubsub := dbpubsub.NewInMemory()
|
||||
clock := quartz.NewReal()
|
||||
|
||||
ownerID := uuid.New()
|
||||
chatID := uuid.New()
|
||||
modelConfigID := uuid.New()
|
||||
userPrompt := "review pull request 23633 and fix review threads"
|
||||
wantTitle := "Review PR 23633"
|
||||
|
||||
chat := database.Chat{
|
||||
ID: chatID,
|
||||
OwnerID: ownerID,
|
||||
LastModelConfigID: modelConfigID,
|
||||
Status: database.ChatStatusCompleted,
|
||||
Title: fallbackChatTitle(userPrompt),
|
||||
}
|
||||
modelConfig := database.ChatModelConfig{
|
||||
ID: modelConfigID,
|
||||
Provider: "anthropic",
|
||||
Model: "claude-haiku-4-5",
|
||||
ContextLimit: 8192,
|
||||
}
|
||||
updatedChat := chat
|
||||
updatedChat.Title = wantTitle
|
||||
|
||||
messageEvents := make(chan struct {
|
||||
payload coderdpubsub.ChatEvent
|
||||
err error
|
||||
}, 1)
|
||||
cancelSub, err := pubsub.SubscribeWithErr(
|
||||
coderdpubsub.ChatEventChannel(ownerID),
|
||||
coderdpubsub.HandleChatEvent(func(_ context.Context, payload coderdpubsub.ChatEvent, err error) {
|
||||
messageEvents <- struct {
|
||||
payload coderdpubsub.ChatEvent
|
||||
err error
|
||||
}{payload: payload, err: err}
|
||||
}),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
defer cancelSub()
|
||||
|
||||
serverURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse {
|
||||
require.Equal(t, "claude-haiku-4-5", req.Model)
|
||||
return chattest.AnthropicNonStreamingResponse(wantTitle)
|
||||
})
|
||||
|
||||
server := &Server{
|
||||
db: db,
|
||||
logger: logger,
|
||||
pubsub: pubsub,
|
||||
configCache: newChatConfigCache(context.Background(), db, clock),
|
||||
}
|
||||
|
||||
db.EXPECT().GetChatModelConfigByID(gomock.Any(), modelConfigID).Return(modelConfig, nil)
|
||||
db.EXPECT().GetEnabledChatProviders(gomock.Any()).Return([]database.ChatProvider{{
|
||||
Provider: "anthropic",
|
||||
APIKey: "test-key",
|
||||
BaseUrl: serverURL,
|
||||
}}, nil)
|
||||
db.EXPECT().GetChatUsageLimitConfig(gomock.Any()).Return(database.ChatUsageLimitConfig{}, sql.ErrNoRows)
|
||||
db.EXPECT().GetChatMessagesByChatIDAscPaginated(
|
||||
gomock.Any(),
|
||||
database.GetChatMessagesByChatIDAscPaginatedParams{
|
||||
ChatID: chatID,
|
||||
AfterID: 0,
|
||||
LimitVal: manualTitleMessageWindowLimit,
|
||||
},
|
||||
).Return([]database.ChatMessage{
|
||||
mustChatMessage(
|
||||
t,
|
||||
database.ChatMessageRoleUser,
|
||||
database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessageText(userPrompt),
|
||||
),
|
||||
mustChatMessage(
|
||||
t,
|
||||
database.ChatMessageRoleAssistant,
|
||||
database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessageText("checking the diff now"),
|
||||
),
|
||||
}, nil)
|
||||
db.EXPECT().GetChatMessagesByChatIDDescPaginated(
|
||||
gomock.Any(),
|
||||
database.GetChatMessagesByChatIDDescPaginatedParams{
|
||||
ChatID: chatID,
|
||||
BeforeID: 0,
|
||||
LimitVal: manualTitleMessageWindowLimit,
|
||||
},
|
||||
).Return(nil, nil)
|
||||
db.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return(nil, nil)
|
||||
|
||||
gomock.InOrder(
|
||||
db.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("chat_title_regenerate_lock")).DoAndReturn(
|
||||
func(fn func(database.Store) error, opts *database.TxOptions) error {
|
||||
require.Equal(t, "chat_title_regenerate_lock", opts.TxIdentifier)
|
||||
return fn(lockTx)
|
||||
},
|
||||
),
|
||||
db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn(
|
||||
func(fn func(database.Store) error, opts *database.TxOptions) error {
|
||||
require.Nil(t, opts)
|
||||
return fn(usageTx)
|
||||
},
|
||||
),
|
||||
db.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("chat_title_regenerate_unlock")).DoAndReturn(
|
||||
func(fn func(database.Store) error, opts *database.TxOptions) error {
|
||||
require.Equal(t, "chat_title_regenerate_unlock", opts.TxIdentifier)
|
||||
return fn(unlockTx)
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
lockTx.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return(chat, nil)
|
||||
lockTx.EXPECT().UpdateChatStatusPreserveUpdatedAt(gomock.Any(), gomock.AssignableToTypeOf(database.UpdateChatStatusPreserveUpdatedAtParams{})).DoAndReturn(
|
||||
func(_ context.Context, arg database.UpdateChatStatusPreserveUpdatedAtParams) (database.Chat, error) {
|
||||
require.Equal(t, chatID, arg.ID)
|
||||
require.Equal(t, chat.Status, arg.Status)
|
||||
require.Equal(t, uuid.NullUUID{UUID: manualTitleLockWorkerID, Valid: true}, arg.WorkerID)
|
||||
require.True(t, arg.StartedAt.Valid)
|
||||
require.True(t, arg.HeartbeatAt.Valid)
|
||||
return chat, nil
|
||||
},
|
||||
)
|
||||
|
||||
usageTx.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return(chat, nil)
|
||||
usageTx.EXPECT().InsertChatMessages(gomock.Any(), gomock.AssignableToTypeOf(database.InsertChatMessagesParams{})).DoAndReturn(
|
||||
func(_ context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) {
|
||||
require.Equal(t, []uuid.UUID{ownerID}, arg.CreatedBy)
|
||||
require.Equal(t, []uuid.UUID{modelConfigID}, arg.ModelConfigID)
|
||||
require.Equal(t, []string{"[]"}, arg.Content)
|
||||
return []database.ChatMessage{{ID: 91}}, nil
|
||||
},
|
||||
)
|
||||
usageTx.EXPECT().SoftDeleteChatMessageByID(gomock.Any(), int64(91)).Return(nil)
|
||||
usageTx.EXPECT().UpdateChatByID(gomock.Any(), database.UpdateChatByIDParams{
|
||||
ID: chatID,
|
||||
Title: wantTitle,
|
||||
}).Return(updatedChat, nil)
|
||||
|
||||
lockedChatWithMarker := updatedChat
|
||||
lockedChatWithMarker.WorkerID = uuid.NullUUID{UUID: manualTitleLockWorkerID, Valid: true}
|
||||
unlockTx.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return(lockedChatWithMarker, nil)
|
||||
unlockTx.EXPECT().UpdateChatStatusPreserveUpdatedAt(gomock.Any(), gomock.AssignableToTypeOf(database.UpdateChatStatusPreserveUpdatedAtParams{})).DoAndReturn(
|
||||
func(_ context.Context, arg database.UpdateChatStatusPreserveUpdatedAtParams) (database.Chat, error) {
|
||||
require.Equal(t, chatID, arg.ID)
|
||||
require.False(t, arg.WorkerID.Valid)
|
||||
require.False(t, arg.StartedAt.Valid)
|
||||
require.False(t, arg.HeartbeatAt.Valid)
|
||||
return updatedChat, nil
|
||||
},
|
||||
)
|
||||
|
||||
gotChat, err := server.RegenerateChatTitle(ctx, chat)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, updatedChat, gotChat)
|
||||
|
||||
select {
|
||||
case event := <-messageEvents:
|
||||
require.NoError(t, event.err)
|
||||
require.Equal(t, coderdpubsub.ChatEventKindTitleChange, event.payload.Kind)
|
||||
require.Equal(t, chatID, event.payload.Chat.ID)
|
||||
require.Equal(t, wantTitle, event.payload.Chat.Title)
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("timed out waiting for title change pubsub event")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRefreshChatWorkspaceSnapshot_NoReloadWhenWorkspacePresent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
+229
-12
@@ -2,6 +2,8 @@ package chatd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -36,6 +38,17 @@ const titleGenerationPrompt = "You are a title generator. Your ONLY job is to ou
|
||||
"Output ONLY the title — no quotes, no emoji, no markdown, no code fences, " +
|
||||
"no trailing punctuation, no preamble, no explanation. Sentence case."
|
||||
|
||||
const (
|
||||
// maxConversationContextRunes caps the conversation sample in manual
|
||||
// title prompts to avoid exceeding model context windows.
|
||||
maxConversationContextRunes = 6000
|
||||
// maxLatestUserMessageRunes caps the latest user message excerpt.
|
||||
maxLatestUserMessageRunes = 1000
|
||||
// recentTurnWindow is the number of most recent turns included
|
||||
// alongside the first user turn in manual title context.
|
||||
recentTurnWindow = 3
|
||||
)
|
||||
|
||||
// preferredTitleModels are lightweight models used for title
|
||||
// generation, one per provider type. Each entry uses the
|
||||
// cheapest/fastest small model for that provider as identified
|
||||
@@ -54,6 +67,33 @@ var preferredTitleModels = []struct {
|
||||
{fantasyvercel.Name, "anthropic/claude-haiku-4.5"},
|
||||
}
|
||||
|
||||
func selectPreferredConfiguredShortTextModelConfig(
|
||||
configs []database.ChatModelConfig,
|
||||
) (database.ChatModelConfig, bool) {
|
||||
for _, preferred := range preferredTitleModels {
|
||||
for _, config := range configs {
|
||||
if chatprovider.NormalizeProvider(config.Provider) != preferred.provider {
|
||||
continue
|
||||
}
|
||||
if !strings.EqualFold(strings.TrimSpace(config.Model), preferred.model) {
|
||||
continue
|
||||
}
|
||||
return config, true
|
||||
}
|
||||
}
|
||||
return database.ChatModelConfig{}, false
|
||||
}
|
||||
|
||||
func normalizeShortTextOutput(text string) string {
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
text = strings.Trim(text, "\"'`")
|
||||
return strings.Join(strings.Fields(text), " ")
|
||||
}
|
||||
|
||||
// maybeGenerateChatTitle generates an AI title for the chat when
|
||||
// appropriate (first user message, no assistant reply yet, and the
|
||||
// current title is either empty or still the fallback truncation).
|
||||
@@ -138,7 +178,7 @@ func generateTitle(
|
||||
model fantasy.LanguageModel,
|
||||
input string,
|
||||
) (string, error) {
|
||||
title, err := generateShortText(ctx, model, titleGenerationPrompt, input)
|
||||
title, _, err := generateShortText(ctx, model, titleGenerationPrompt, input)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
@@ -199,13 +239,10 @@ func titleInput(
|
||||
}
|
||||
|
||||
func normalizeTitleOutput(title string) string {
|
||||
title = strings.TrimSpace(title)
|
||||
title = normalizeShortTextOutput(title)
|
||||
if title == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
title = strings.Trim(title, "\"'`")
|
||||
title = strings.Join(strings.Fields(title), " ")
|
||||
return truncateRunes(title, 80)
|
||||
}
|
||||
|
||||
@@ -226,7 +263,7 @@ func fallbackChatTitle(message string) string {
|
||||
|
||||
title := strings.Join(words, " ")
|
||||
if truncated {
|
||||
title += "…"
|
||||
return truncateRunes(title, maxRunes-1) + "…"
|
||||
}
|
||||
|
||||
return truncateRunes(title, maxRunes)
|
||||
@@ -260,6 +297,186 @@ func truncateRunes(value string, maxLen int) string {
|
||||
return string(runes[:maxLen])
|
||||
}
|
||||
|
||||
// Manual title regeneration is user-initiated and can use richer
|
||||
// conversation context than the automatic first-message title path
|
||||
// above. These helpers keep the manual prompt-building logic private
|
||||
// while reusing the shared title-generation utilities in this file.
|
||||
type manualTitleTurn struct {
|
||||
role string
|
||||
text string
|
||||
}
|
||||
|
||||
func extractManualTitleTurns(messages []database.ChatMessage) []manualTitleTurn {
|
||||
turns := make([]manualTitleTurn, 0, len(messages))
|
||||
for _, message := range messages {
|
||||
if message.Visibility == database.ChatMessageVisibilityModel {
|
||||
continue
|
||||
}
|
||||
|
||||
role := ""
|
||||
switch message.Role {
|
||||
case database.ChatMessageRoleUser:
|
||||
role = string(database.ChatMessageRoleUser)
|
||||
case database.ChatMessageRoleAssistant:
|
||||
role = string(database.ChatMessageRoleAssistant)
|
||||
default:
|
||||
continue
|
||||
}
|
||||
|
||||
parts, err := chatprompt.ParseContent(message)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
text := strings.TrimSpace(contentBlocksToText(parts))
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
turns = append(turns, manualTitleTurn{
|
||||
role: role,
|
||||
text: text,
|
||||
})
|
||||
}
|
||||
|
||||
return turns
|
||||
}
|
||||
|
||||
func selectManualTitleTurnIndexes(turns []manualTitleTurn) []int {
|
||||
firstUserIndex := slices.IndexFunc(turns, func(turn manualTitleTurn) bool {
|
||||
return turn.role == string(database.ChatMessageRoleUser)
|
||||
})
|
||||
if firstUserIndex == -1 {
|
||||
return nil
|
||||
}
|
||||
|
||||
windowStart := max(0, len(turns)-recentTurnWindow)
|
||||
selected := make([]int, 0, recentTurnWindow+1)
|
||||
if firstUserIndex < windowStart {
|
||||
selected = append(selected, firstUserIndex)
|
||||
}
|
||||
for i := windowStart; i < len(turns); i++ {
|
||||
selected = append(selected, i)
|
||||
}
|
||||
|
||||
return selected
|
||||
}
|
||||
|
||||
func buildManualTitleContext(
|
||||
turns []manualTitleTurn,
|
||||
selected []int,
|
||||
) (conversationBlock string, latestUserMsg string) {
|
||||
userCount := 0
|
||||
for _, turn := range turns {
|
||||
if turn.role != string(database.ChatMessageRoleUser) {
|
||||
continue
|
||||
}
|
||||
userCount++
|
||||
latestUserMsg = turn.text
|
||||
}
|
||||
|
||||
latestUserMsg = truncateRunes(latestUserMsg, maxLatestUserMessageRunes)
|
||||
if userCount <= 1 || len(selected) == 0 {
|
||||
return "", latestUserMsg
|
||||
}
|
||||
|
||||
lines := make([]string, 0, len(selected)+1)
|
||||
for i, idx := range selected {
|
||||
if i == 1 {
|
||||
if gap := idx - selected[i-1] - 1; gap > 0 {
|
||||
lines = append(lines, fmt.Sprintf("[... %d earlier turns omitted ...]", gap))
|
||||
}
|
||||
}
|
||||
lines = append(lines, fmt.Sprintf("[%s]: %s", turns[idx].role, turns[idx].text))
|
||||
}
|
||||
|
||||
conversationBlock = strings.Join(lines, "\n")
|
||||
conversationBlock = truncateRunes(conversationBlock, maxConversationContextRunes)
|
||||
return conversationBlock, latestUserMsg
|
||||
}
|
||||
|
||||
func renderManualTitlePrompt(
|
||||
conversationBlock string,
|
||||
firstUserText string,
|
||||
latestUserMsg string,
|
||||
) string {
|
||||
var prompt strings.Builder
|
||||
write := func(value string) {
|
||||
_, _ = prompt.WriteString(value)
|
||||
}
|
||||
|
||||
write("You are a title generator for an AI coding assistant conversation.\n\n")
|
||||
write("The user's primary objective was:\n<primary_objective>\n")
|
||||
write(firstUserText)
|
||||
write("\n</primary_objective>")
|
||||
|
||||
if conversationBlock != "" {
|
||||
write("\n\nConversation sample:\n<conversation_sample>\n")
|
||||
write(conversationBlock)
|
||||
write("\n</conversation_sample>")
|
||||
}
|
||||
|
||||
if strings.TrimSpace(latestUserMsg) != strings.TrimSpace(truncateRunes(firstUserText, maxLatestUserMessageRunes)) {
|
||||
write("\n\nThe user's most recent message:\n<latest_message>\n")
|
||||
write(latestUserMsg)
|
||||
write("\n</latest_message>\n")
|
||||
write("Note: Weight the overall conversation arc more heavily than just the latest message.")
|
||||
}
|
||||
|
||||
write("\n\nRequirements:\n")
|
||||
write("- Output a short title of 2-8 words.\n")
|
||||
write("- Use verb-noun format in sentence case.\n")
|
||||
write("- Preserve specific identifiers (PR numbers, repo names, file paths, function names, error messages).\n")
|
||||
write("- No trailing punctuation, quotes, emoji, or markdown.\n")
|
||||
write("- No temporal phrasing (\"Continue\", \"Follow up on\") or meta phrasing (\"Chat about\").\n")
|
||||
write("- Output ONLY the title - nothing else.\n")
|
||||
return prompt.String()
|
||||
}
|
||||
|
||||
func generateManualTitle(
|
||||
ctx context.Context,
|
||||
messages []database.ChatMessage,
|
||||
fallbackModel fantasy.LanguageModel,
|
||||
) (string, fantasy.Usage, error) {
|
||||
turns := extractManualTitleTurns(messages)
|
||||
selected := selectManualTitleTurnIndexes(turns)
|
||||
|
||||
firstUserIndex := slices.IndexFunc(turns, func(turn manualTitleTurn) bool {
|
||||
return turn.role == string(database.ChatMessageRoleUser)
|
||||
})
|
||||
if firstUserIndex == -1 {
|
||||
return "", fantasy.Usage{}, nil
|
||||
}
|
||||
firstUserText := truncateRunes(turns[firstUserIndex].text, maxLatestUserMessageRunes)
|
||||
|
||||
conversationBlock, latestUserMsg := buildManualTitleContext(turns, selected)
|
||||
systemPrompt := renderManualTitlePrompt(
|
||||
conversationBlock,
|
||||
firstUserText,
|
||||
latestUserMsg,
|
||||
)
|
||||
|
||||
titleCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
title, usage, err := generateShortText(
|
||||
titleCtx,
|
||||
fallbackModel,
|
||||
systemPrompt,
|
||||
"Generate the title.",
|
||||
)
|
||||
if err != nil {
|
||||
return "", fantasy.Usage{}, err
|
||||
}
|
||||
|
||||
title = normalizeTitleOutput(title)
|
||||
if title == "" {
|
||||
return "", usage, xerrors.New("generated title was empty")
|
||||
}
|
||||
|
||||
return title, usage, nil
|
||||
}
|
||||
|
||||
const pushSummaryPrompt = "You are a notification assistant. Given a chat title " +
|
||||
"and the agent's last message, write a single short sentence (under 100 characters) " +
|
||||
"summarizing what the agent did. This will be shown as a push notification body. " +
|
||||
@@ -281,6 +498,7 @@ func generatePushSummary(
|
||||
summaryCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
assistantText = truncateRunes(assistantText, maxConversationContextRunes)
|
||||
input := "Chat title: " + chat.Title + "\n\nAgent's last message:\n" + assistantText
|
||||
|
||||
candidates := make([]fantasy.LanguageModel, 0, len(preferredTitleModels)+1)
|
||||
@@ -296,7 +514,7 @@ func generatePushSummary(
|
||||
candidates = append(candidates, fallbackModel)
|
||||
|
||||
for _, model := range candidates {
|
||||
summary, err := generateShortText(summaryCtx, model, pushSummaryPrompt, input)
|
||||
summary, _, err := generateShortText(summaryCtx, model, pushSummaryPrompt, input)
|
||||
if err != nil {
|
||||
logger.Debug(ctx, "push summary model candidate failed",
|
||||
slog.Error(err),
|
||||
@@ -318,7 +536,7 @@ func generateShortText(
|
||||
model fantasy.LanguageModel,
|
||||
systemPrompt string,
|
||||
userInput string,
|
||||
) (string, error) {
|
||||
) (string, fantasy.Usage, error) {
|
||||
prompt := []fantasy.Message{
|
||||
{
|
||||
Role: fantasy.MessageRoleSystem,
|
||||
@@ -346,7 +564,7 @@ func generateShortText(
|
||||
return genErr
|
||||
}, nil)
|
||||
if err != nil {
|
||||
return "", xerrors.Errorf("generate short text: %w", err)
|
||||
return "", fantasy.Usage{}, xerrors.Errorf("generate short text: %w", err)
|
||||
}
|
||||
|
||||
responseParts := make([]codersdk.ChatMessagePart, 0, len(response.Content))
|
||||
@@ -355,7 +573,6 @@ func generateShortText(
|
||||
responseParts = append(responseParts, p)
|
||||
}
|
||||
}
|
||||
text := strings.TrimSpace(contentBlocksToText(responseParts))
|
||||
text = strings.Trim(text, "\"'`")
|
||||
return text, nil
|
||||
text := normalizeShortTextOutput(contentBlocksToText(responseParts))
|
||||
return text, response.Usage, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,584 @@
|
||||
package chatd //nolint:testpackage // Keeps internal helper tests in-package.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"github.com/sqlc-dev/pqtype"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
)
|
||||
|
||||
func Test_extractManualTitleTurns(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
messages []database.ChatMessage
|
||||
want []manualTitleTurn
|
||||
}{
|
||||
{
|
||||
name: "filters to visible user and assistant text turns",
|
||||
messages: []database.ChatMessage{
|
||||
mustChatMessage(t, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: " review quickgen helpers "},
|
||||
),
|
||||
mustChatMessage(t, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: " drafted a plan "},
|
||||
),
|
||||
mustChatMessage(t, database.ChatMessageRoleSystem, database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "system prompt"},
|
||||
),
|
||||
mustChatMessage(t, database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "tool output"},
|
||||
),
|
||||
mustChatMessage(t, database.ChatMessageRoleUser, database.ChatMessageVisibilityModel,
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "hidden model note"},
|
||||
),
|
||||
mustChatMessage(t, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: " "},
|
||||
),
|
||||
mustChatMessage(t, database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeReasoning, Text: "reasoning only"},
|
||||
),
|
||||
mustChatMessage(t, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeFile, MediaType: "text/plain"},
|
||||
),
|
||||
},
|
||||
want: []manualTitleTurn{
|
||||
{role: "user", text: "review quickgen helpers"},
|
||||
{role: "assistant", text: "drafted a plan"},
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "reuses text extraction for multi-part content",
|
||||
messages: []database.ChatMessage{
|
||||
mustChatMessage(t, database.ChatMessageRoleUser, database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: "first chunk"},
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeReasoning, Text: "skip me"},
|
||||
codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeText, Text: " second chunk "},
|
||||
),
|
||||
},
|
||||
want: []manualTitleTurn{{role: "user", text: "first chunk second chunk"}},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got := extractManualTitleTurns(tt.messages)
|
||||
require.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func Test_selectManualTitleTurnIndexes(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
turns []manualTitleTurn
|
||||
want []int
|
||||
}{
|
||||
{
|
||||
name: "single user turn",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "user", text: "one"},
|
||||
},
|
||||
want: []int{0},
|
||||
},
|
||||
{
|
||||
name: "first user plus trailing window",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "user", text: "one"},
|
||||
{role: "assistant", text: "two"},
|
||||
{role: "user", text: "three"},
|
||||
{role: "assistant", text: "four"},
|
||||
{role: "user", text: "five"},
|
||||
},
|
||||
want: []int{0, 2, 3, 4},
|
||||
},
|
||||
{
|
||||
name: "two turns returns both",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "user", text: "one"},
|
||||
{role: "assistant", text: "two"},
|
||||
},
|
||||
want: []int{0, 1},
|
||||
},
|
||||
{
|
||||
name: "prepends first user when before trailing window",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "assistant", text: "intro"},
|
||||
{role: "assistant", text: "setup"},
|
||||
{role: "user", text: "goal"},
|
||||
{role: "assistant", text: "a"},
|
||||
{role: "assistant", text: "b"},
|
||||
{role: "assistant", text: "c"},
|
||||
},
|
||||
want: []int{2, 3, 4, 5},
|
||||
},
|
||||
{
|
||||
name: "ten plus turns keeps first user and last three",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "assistant", text: "0"},
|
||||
{role: "assistant", text: "1"},
|
||||
{role: "user", text: "2"},
|
||||
{role: "assistant", text: "3"},
|
||||
{role: "assistant", text: "4"},
|
||||
{role: "assistant", text: "5"},
|
||||
{role: "assistant", text: "6"},
|
||||
{role: "assistant", text: "7"},
|
||||
{role: "assistant", text: "8"},
|
||||
{role: "user", text: "9"},
|
||||
{role: "assistant", text: "10"},
|
||||
{role: "user", text: "11"},
|
||||
},
|
||||
want: []int{2, 9, 10, 11},
|
||||
},
|
||||
{
|
||||
name: "no user turns",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "assistant", text: "one"},
|
||||
{role: "assistant", text: "two"},
|
||||
},
|
||||
want: nil,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got := selectManualTitleTurnIndexes(tt.turns)
|
||||
require.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func Test_buildManualTitleContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
longConversationText := strings.Repeat("a", 3500)
|
||||
longLatestUserText := strings.Repeat("z", 1200)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
turns []manualTitleTurn
|
||||
selected []int
|
||||
wantConversation string
|
||||
wantConversationEmpty bool
|
||||
wantConversationHasGap bool
|
||||
wantConversationRunes int
|
||||
wantLatestUser string
|
||||
wantLatestUserRunes int
|
||||
wantLatestUserContains string
|
||||
wantLatestUserNotEmpty bool
|
||||
}{
|
||||
{
|
||||
name: "adds gap marker when selected turns skip earlier context",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "user", text: "open pull request"},
|
||||
{role: "assistant", text: "checked CI"},
|
||||
{role: "user", text: "review logs"},
|
||||
{role: "assistant", text: "found flaky test"},
|
||||
{role: "user", text: "update chat title"},
|
||||
},
|
||||
selected: []int{0, 3, 4},
|
||||
wantConversationHasGap: true,
|
||||
wantLatestUser: "update chat title",
|
||||
},
|
||||
{
|
||||
name: "omits gap marker for contiguous selection",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "user", text: "open pull request"},
|
||||
{role: "assistant", text: "checked CI"},
|
||||
{role: "user", text: "update chat title"},
|
||||
},
|
||||
selected: []int{0, 1, 2},
|
||||
wantConversation: "[user]: open pull request\n[assistant]: checked CI\n[user]: update chat title",
|
||||
wantConversationHasGap: false,
|
||||
wantLatestUser: "update chat title",
|
||||
},
|
||||
{
|
||||
name: "single useful user turn returns empty conversation block",
|
||||
turns: []manualTitleTurn{{role: "user", text: "rename helper"}},
|
||||
selected: []int{0},
|
||||
wantConversationEmpty: true,
|
||||
wantLatestUser: "rename helper",
|
||||
},
|
||||
{
|
||||
name: "truncates conversation block at six thousand runes",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "user", text: longConversationText},
|
||||
{role: "assistant", text: longConversationText},
|
||||
{role: "user", text: "latest"},
|
||||
},
|
||||
selected: []int{0, 1, 2},
|
||||
wantConversationRunes: 6000,
|
||||
wantLatestUser: "latest",
|
||||
},
|
||||
{
|
||||
name: "truncates latest user message at one thousand runes",
|
||||
turns: []manualTitleTurn{
|
||||
{role: "user", text: "first"},
|
||||
{role: "assistant", text: "reply"},
|
||||
{role: "user", text: longLatestUserText},
|
||||
},
|
||||
selected: []int{0, 1, 2},
|
||||
wantLatestUserRunes: 1000,
|
||||
wantLatestUserContains: strings.Repeat("z", 1000),
|
||||
wantLatestUserNotEmpty: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
conversationBlock, latestUserMsg := buildManualTitleContext(tt.turns, tt.selected)
|
||||
|
||||
if tt.wantConversationEmpty {
|
||||
require.Empty(t, conversationBlock)
|
||||
}
|
||||
if tt.wantConversation != "" {
|
||||
require.Equal(t, tt.wantConversation, conversationBlock)
|
||||
}
|
||||
if tt.wantConversationHasGap {
|
||||
require.Contains(t, conversationBlock, "[... 2 earlier turns omitted ...]")
|
||||
} else if !tt.wantConversationEmpty {
|
||||
require.NotContains(t, conversationBlock, "earlier turns omitted")
|
||||
}
|
||||
if tt.wantConversationRunes > 0 {
|
||||
require.Len(t, []rune(conversationBlock), tt.wantConversationRunes)
|
||||
}
|
||||
if tt.wantLatestUser != "" {
|
||||
require.Equal(t, tt.wantLatestUser, latestUserMsg)
|
||||
}
|
||||
if tt.wantLatestUserRunes > 0 {
|
||||
require.Len(t, []rune(latestUserMsg), tt.wantLatestUserRunes)
|
||||
}
|
||||
if tt.wantLatestUserContains != "" {
|
||||
require.Equal(t, tt.wantLatestUserContains, latestUserMsg)
|
||||
}
|
||||
if tt.wantLatestUserNotEmpty {
|
||||
require.NotEmpty(t, latestUserMsg)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func Test_renderManualTitlePrompt(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
longFirstUserText := strings.Repeat("b", 1501)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
conversationBlock string
|
||||
firstUserText string
|
||||
latestUserMsg string
|
||||
wantConversationSample bool
|
||||
wantLatestSection bool
|
||||
}{
|
||||
{
|
||||
name: "includes conversation sample when provided",
|
||||
conversationBlock: "[user]: inspect logs\n[assistant]: found flaky test",
|
||||
firstUserText: "inspect logs",
|
||||
latestUserMsg: "update quickgen title",
|
||||
wantConversationSample: true,
|
||||
wantLatestSection: true,
|
||||
},
|
||||
{
|
||||
name: "omits optional sections when not needed",
|
||||
conversationBlock: "",
|
||||
firstUserText: "inspect logs",
|
||||
latestUserMsg: "inspect logs",
|
||||
wantConversationSample: false,
|
||||
wantLatestSection: false,
|
||||
},
|
||||
{
|
||||
name: "latest section compares trimmed text",
|
||||
conversationBlock: "",
|
||||
firstUserText: "inspect logs",
|
||||
latestUserMsg: " inspect logs ",
|
||||
wantConversationSample: false,
|
||||
wantLatestSection: false,
|
||||
},
|
||||
{
|
||||
name: "omits latest section when same message truncated",
|
||||
conversationBlock: "",
|
||||
firstUserText: longFirstUserText,
|
||||
latestUserMsg: truncateRunes(longFirstUserText, 1000),
|
||||
wantConversationSample: false,
|
||||
wantLatestSection: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
prompt := renderManualTitlePrompt(tt.conversationBlock, tt.firstUserText, tt.latestUserMsg)
|
||||
|
||||
require.Contains(t, prompt, "The user's primary objective was:")
|
||||
require.Contains(t, prompt, "Requirements:")
|
||||
require.Contains(t, prompt, "- Output a short title of 2-8 words.")
|
||||
require.Contains(t, prompt, "- Output ONLY the title - nothing else.")
|
||||
|
||||
if tt.wantConversationSample {
|
||||
require.Contains(t, prompt, "Conversation sample:")
|
||||
require.Contains(t, prompt, tt.conversationBlock)
|
||||
} else {
|
||||
require.NotContains(t, prompt, "Conversation sample:")
|
||||
}
|
||||
|
||||
if tt.wantLatestSection {
|
||||
require.Contains(t, prompt, "The user's most recent message:")
|
||||
require.Contains(t, prompt, "Note: Weight the overall conversation arc more heavily than just the latest message.")
|
||||
require.Contains(t, prompt, strings.TrimSpace(tt.latestUserMsg))
|
||||
} else {
|
||||
require.NotContains(t, prompt, "The user's most recent message:")
|
||||
require.NotContains(t, prompt, "Weight the overall conversation arc more heavily")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func Test_generateManualTitle_UsesTimeout(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
messages := []database.ChatMessage{
|
||||
mustChatMessage(
|
||||
t,
|
||||
database.ChatMessageRoleUser,
|
||||
database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessageText("refresh chat title"),
|
||||
),
|
||||
}
|
||||
|
||||
model := &stubModel{
|
||||
generateFn: func(ctx context.Context, call fantasy.Call) (*fantasy.Response, error) {
|
||||
deadline, ok := ctx.Deadline()
|
||||
require.True(t, ok, "manual title generation should set a deadline")
|
||||
require.WithinDuration(
|
||||
t,
|
||||
time.Now().Add(30*time.Second),
|
||||
deadline,
|
||||
2*time.Second,
|
||||
)
|
||||
require.Len(t, call.Prompt, 2)
|
||||
return &fantasy.Response{
|
||||
Content: fantasy.ResponseContent{
|
||||
fantasy.TextContent{Text: "Refresh title"},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
title, _, err := generateManualTitle(
|
||||
context.Background(),
|
||||
messages,
|
||||
model,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "Refresh title", title)
|
||||
}
|
||||
|
||||
func Test_generateManualTitle_TruncatesFirstUserInput(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
longFirstUserText := strings.Repeat("a", 1500)
|
||||
messages := []database.ChatMessage{
|
||||
mustChatMessage(
|
||||
t,
|
||||
database.ChatMessageRoleUser,
|
||||
database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessageText(longFirstUserText),
|
||||
),
|
||||
}
|
||||
|
||||
model := &stubModel{
|
||||
generateFn: func(_ context.Context, call fantasy.Call) (*fantasy.Response, error) {
|
||||
require.Len(t, call.Prompt, 2)
|
||||
systemText, ok := call.Prompt[0].Content[0].(fantasy.TextPart)
|
||||
require.True(t, ok)
|
||||
require.Contains(t, systemText.Text, truncateRunes(longFirstUserText, 1000))
|
||||
|
||||
userText, ok := call.Prompt[1].Content[0].(fantasy.TextPart)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "Generate the title.", userText.Text)
|
||||
return &fantasy.Response{
|
||||
Content: fantasy.ResponseContent{
|
||||
fantasy.TextContent{Text: "Refresh title"},
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
_, _, err := generateManualTitle(
|
||||
context.Background(),
|
||||
messages,
|
||||
model,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
func Test_generateManualTitle_ReturnsUsageForEmptyNormalizedTitle(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
messages := []database.ChatMessage{
|
||||
mustChatMessage(
|
||||
t,
|
||||
database.ChatMessageRoleUser,
|
||||
database.ChatMessageVisibilityBoth,
|
||||
codersdk.ChatMessageText("refresh chat title"),
|
||||
),
|
||||
}
|
||||
|
||||
model := &stubModel{
|
||||
generateFn: func(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) {
|
||||
return &fantasy.Response{
|
||||
Content: fantasy.ResponseContent{
|
||||
fantasy.TextContent{Text: "\"\""},
|
||||
},
|
||||
Usage: fantasy.Usage{
|
||||
InputTokens: 11,
|
||||
OutputTokens: 7,
|
||||
TotalTokens: 18,
|
||||
},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
_, usage, err := generateManualTitle(
|
||||
context.Background(),
|
||||
messages,
|
||||
model,
|
||||
)
|
||||
require.ErrorContains(t, err, "generated title was empty")
|
||||
require.Equal(t, int64(11), usage.InputTokens)
|
||||
require.Equal(t, int64(7), usage.OutputTokens)
|
||||
require.Equal(t, int64(18), usage.TotalTokens)
|
||||
}
|
||||
|
||||
func Test_selectPreferredConfiguredShortTextModelConfig(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("chooses the highest-priority configured lightweight model", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
configs := []database.ChatModelConfig{
|
||||
{Provider: preferredTitleModels[2].provider, Model: preferredTitleModels[2].model},
|
||||
{Provider: preferredTitleModels[1].provider, Model: preferredTitleModels[1].model},
|
||||
{Provider: "openai", Model: "gpt-4.1"},
|
||||
}
|
||||
|
||||
got, ok := selectPreferredConfiguredShortTextModelConfig(configs)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, preferredTitleModels[1].provider, got.Provider)
|
||||
require.Equal(t, preferredTitleModels[1].model, got.Model)
|
||||
})
|
||||
|
||||
t.Run("returns false when no preferred lightweight model is configured", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
got, ok := selectPreferredConfiguredShortTextModelConfig([]database.ChatModelConfig{{
|
||||
Provider: "openai",
|
||||
Model: "gpt-4.1",
|
||||
}})
|
||||
require.False(t, ok)
|
||||
require.Equal(t, database.ChatModelConfig{}, got)
|
||||
})
|
||||
}
|
||||
|
||||
func Test_generateShortText_NormalizesQuotedOutput(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
model := &stubModel{
|
||||
generateFn: func(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) {
|
||||
return &fantasy.Response{
|
||||
Content: fantasy.ResponseContent{
|
||||
fantasy.TextContent{Text: " \"Quoted summary\" "},
|
||||
},
|
||||
Usage: fantasy.Usage{InputTokens: 3, OutputTokens: 2, TotalTokens: 5},
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
text, usage, err := generateShortText(context.Background(), model, "system", "user")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "Quoted summary", text)
|
||||
require.Equal(t, int64(3), usage.InputTokens)
|
||||
require.Equal(t, int64(2), usage.OutputTokens)
|
||||
require.Equal(t, int64(5), usage.TotalTokens)
|
||||
}
|
||||
|
||||
type stubModel struct {
|
||||
generateFn func(context.Context, fantasy.Call) (*fantasy.Response, error)
|
||||
}
|
||||
|
||||
func (m *stubModel) Generate(
|
||||
ctx context.Context,
|
||||
call fantasy.Call,
|
||||
) (*fantasy.Response, error) {
|
||||
return m.generateFn(ctx, call)
|
||||
}
|
||||
|
||||
func (*stubModel) Stream(
|
||||
context.Context,
|
||||
fantasy.Call,
|
||||
) (fantasy.StreamResponse, error) {
|
||||
return nil, xerrors.New("stream not implemented")
|
||||
}
|
||||
|
||||
func (*stubModel) GenerateObject(
|
||||
context.Context,
|
||||
fantasy.ObjectCall,
|
||||
) (*fantasy.ObjectResponse, error) {
|
||||
return nil, xerrors.New("generate object not implemented")
|
||||
}
|
||||
|
||||
func (*stubModel) StreamObject(
|
||||
context.Context,
|
||||
fantasy.ObjectCall,
|
||||
) (fantasy.ObjectStreamResponse, error) {
|
||||
return nil, xerrors.New("stream object not implemented")
|
||||
}
|
||||
|
||||
func (*stubModel) Provider() string {
|
||||
return "test"
|
||||
}
|
||||
|
||||
func (*stubModel) Model() string {
|
||||
return "test"
|
||||
}
|
||||
|
||||
func mustChatMessage(
|
||||
t *testing.T,
|
||||
role database.ChatMessageRole,
|
||||
visibility database.ChatMessageVisibility,
|
||||
parts ...codersdk.ChatMessagePart,
|
||||
) database.ChatMessage {
|
||||
t.Helper()
|
||||
|
||||
content, err := json.Marshal(parts)
|
||||
require.NoError(t, err)
|
||||
|
||||
return database.ChatMessage{
|
||||
Role: role,
|
||||
Visibility: visibility,
|
||||
Content: pqtype.NullRawMessage{
|
||||
RawMessage: content,
|
||||
Valid: len(content) > 0,
|
||||
},
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user