Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions coderd/database/dbauthz/dbauthz.go
Original file line number Diff line number Diff line change
Expand Up @@ -3450,6 +3450,13 @@ func (q *querier) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([
return q.db.GetChatStreamSyncRows(ctx, ids)
}

func (q *querier) GetChatSummaryGenerationModelOverride(ctx context.Context) (string, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return "", err
}
return q.db.GetChatSummaryGenerationModelOverride(ctx)
}

func (q *querier) GetChatSystemPrompt(ctx context.Context) (string, error) {
// The system prompt is a deployment-wide setting read during chat
// creation by every authenticated user, so no RBAC policy check
Expand Down Expand Up @@ -8711,6 +8718,13 @@ func (q *querier) UpsertChatRetentionDays(ctx context.Context, retentionDays int
return q.db.UpsertChatRetentionDays(ctx, retentionDays)
}

func (q *querier) UpsertChatSummaryGenerationModelOverride(ctx context.Context, value string) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return err
}
return q.db.UpsertChatSummaryGenerationModelOverride(ctx, value)
}

func (q *querier) UpsertChatSystemPrompt(ctx context.Context, value string) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return err
Expand Down
8 changes: 8 additions & 0 deletions coderd/database/dbauthz/dbauthz_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1201,6 +1201,10 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().GetChatTitleGenerationModelOverride(gomock.Any()).Return("", nil).AnyTimes()
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead)
}))
s.Run("GetChatSummaryGenerationModelOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().GetChatSummaryGenerationModelOverride(gomock.Any()).Return("", nil).AnyTimes()
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead)
}))
s.Run("GetChatPlanModeInstructions", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().GetChatPlanModeInstructions(gomock.Any()).Return("", nil).AnyTimes()
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
Expand Down Expand Up @@ -1661,6 +1665,10 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().UpsertChatTitleGenerationModelOverride(gomock.Any(), "").Return(nil).AnyTimes()
check.Args("").Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
}))
s.Run("UpsertChatSummaryGenerationModelOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().UpsertChatSummaryGenerationModelOverride(gomock.Any(), "").Return(nil).AnyTimes()
check.Args("").Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
}))
s.Run("UpsertChatPlanModeInstructions", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().UpsertChatPlanModeInstructions(gomock.Any(), "").Return(nil).AnyTimes()
check.Args("").Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
Expand Down
16 changes: 16 additions & 0 deletions coderd/database/dbmetrics/querymetrics.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

29 changes: 29 additions & 0 deletions coderd/database/dbmock/dbmock.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions coderd/database/querier.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

22 changes: 22 additions & 0 deletions coderd/database/queries.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 8 additions & 0 deletions coderd/database/queries/siteconfig.sql
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,14 @@ SELECT
INSERT INTO site_configs (key, value) VALUES ('agents_chat_title_generation_model_override', $1)
ON CONFLICT (key) DO UPDATE SET value = $1 WHERE site_configs.key = 'agents_chat_title_generation_model_override';

-- name: GetChatSummaryGenerationModelOverride :one
SELECT
COALESCE((SELECT value FROM site_configs WHERE key = 'agents_chat_summary_generation_model_override'), '') :: text AS model_config_id;

-- name: UpsertChatSummaryGenerationModelOverride :exec
INSERT INTO site_configs (key, value) VALUES ('agents_chat_summary_generation_model_override', $1)
ON CONFLICT (key) DO UPDATE SET value = $1 WHERE site_configs.key = 'agents_chat_summary_generation_model_override';

-- name: GetChatDesktopEnabled :one
SELECT
COALESCE((SELECT value = 'true' FROM site_configs WHERE key = 'agents_desktop_enabled'), false) :: boolean AS enable_desktop;
Expand Down
6 changes: 6 additions & 0 deletions coderd/exp_chats.go
Original file line number Diff line number Diff line change
Expand Up @@ -621,6 +621,12 @@ func (api *API) chatModelOverrideSiteConfig(
getter: api.Database.GetChatTitleGenerationModelOverride,
upsert: api.Database.UpsertChatTitleGenerationModelOverride,
}, nil
case codersdk.ChatModelOverrideContextSummaryGeneration:
return chatModelOverrideSiteConfig{
label: "summary generation",
getter: api.Database.GetChatSummaryGenerationModelOverride,
upsert: api.Database.UpsertChatSummaryGenerationModelOverride,
}, nil
default:
return chatModelOverrideSiteConfig{}, xerrors.Errorf(
"unknown chat model override context %q",
Expand Down
14 changes: 12 additions & 2 deletions coderd/exp_chats_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -11667,6 +11667,16 @@ func TestChatModelOverrides(t *testing.T) {
return db.UpsertChatTitleGenerationModelOverride(dbauthz.AsSystemRestricted(ctx), value)
},
},
{
name: "SummaryGeneration",
context: codersdk.ChatModelOverrideContextSummaryGeneration,
dbGet: func(ctx context.Context, db database.Store) (string, error) {
return db.GetChatSummaryGenerationModelOverride(dbauthz.AsSystemRestricted(ctx))
},
dbUpsert: func(ctx context.Context, db database.Store, value string) error {
return db.UpsertChatSummaryGenerationModelOverride(dbauthz.AsSystemRestricted(ctx), value)
},
},
}

for _, setting := range settings {
Expand Down Expand Up @@ -11812,7 +11822,7 @@ func TestChatModelOverrides(t *testing.T) {
require.Equal(t, "Invalid chat model override context.", sdkErr.Message)
require.Equal(
t,
`Expected one of general, explore, title_generation. Got "not-a-context".`,
`Expected one of general, explore, title_generation, summary_generation. Got "not-a-context".`,
sdkErr.Detail,
)

Expand All @@ -11821,7 +11831,7 @@ func TestChatModelOverrides(t *testing.T) {
require.Equal(t, "Invalid chat model override context.", sdkErr.Message)
require.Equal(
t,
`Expected one of general, explore, title_generation. Got "not-a-context".`,
`Expected one of general, explore, title_generation, summary_generation. Got "not-a-context".`,
sdkErr.Detail,
)
})
Expand Down
29 changes: 27 additions & 2 deletions coderd/x/chatd/chatd.go
Original file line number Diff line number Diff line change
Expand Up @@ -4983,19 +4983,44 @@ func (p *Server) generateAndStoreChatSummary(
p.updateChatSummary(ctx, chat, chat.HistoryVersion, summary, logger)
}

// resolveChatSummaryModel prefers a usable deployment override, skips generation
// on a set-but-unusable one, and otherwise uses the chat's configured model.
func (p *Server) resolveChatSummaryModel(
ctx context.Context,
chat database.Chat,
modelOpts modelBuildOptions,
logger slog.Logger,
) (fantasy.LanguageModel, database.ChatModelConfig, bool) {
//nolint:dogsled // resolveChatModel returns rich routing metadata; summary generation only needs the model and its config.
model, dbConfig, _, _, _, _, _, err := p.resolveChatModel(ctx, chat, modelOpts)
//nolint:dogsled // resolveChatModel returns rich routing metadata; summary generation only needs the model, its config, and provider keys.
model, dbConfig, keys, _, _, _, _, err := p.resolveChatModel(ctx, chat, modelOpts)
if err != nil {
logger.Debug(ctx, "failed to resolve chat model for summary",
slog.F("chat_id", chat.ID), slog.Error(err))
return nil, database.ChatModelConfig{}, false
}

overrideConfig, overrideModel, _, _, overrideSet, overrideErr := p.resolveSummaryGenerationModelOverride(
ctx, chat, keys, modelOpts,
)
if overrideErr != nil {
if overrideSet {
logger.Warn(ctx, "summary generation model override unavailable, skipping summary generation",
slog.F("chat_id", chat.ID),
slog.F("override_context", summaryGenerationOverrideContext),
slog.Error(overrideErr),
)
return nil, database.ChatModelConfig{}, false
}
logger.Debug(ctx, "failed to resolve summary generation model override",
slog.F("chat_id", chat.ID),
slog.F("override_context", summaryGenerationOverrideContext),
slog.Error(overrideErr),
)
}
if overrideSet {
return overrideModel, overrideConfig, true
}

return model, dbConfig, true
}

Expand Down
45 changes: 45 additions & 0 deletions coderd/x/chatd/summary_override.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
package chatd

import (
"context"

"charm.land/fantasy"
"golang.org/x/xerrors"

"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
)

const summaryGenerationOverrideContext = "summary_generation"

func readSummaryGenerationModelOverride(
ctx context.Context,
db database.Store,
) (string, error) {
//nolint:gocritic // Chatd is internal, not a user, so this read uses AsChatd.
chatdCtx := dbauthz.AsChatd(ctx)
raw, err := db.GetChatSummaryGenerationModelOverride(chatdCtx)
if err != nil {
return "", xerrors.Errorf(
"get chat summary generation model override: %w",
err,
)
}
return raw, nil
}

// resolveSummaryGenerationModelOverride resolves the deployment-wide summary
// override. overrideSet reports whether one was configured; if true, any error is
// a hard failure (skip generation), and if false the caller uses the chat's model.
func (p *Server) resolveSummaryGenerationModelOverride(
ctx context.Context,
chat database.Chat,
keys chatprovider.ProviderAPIKeys,
modelOpts modelBuildOptions,
) (database.ChatModelConfig, fantasy.LanguageModel, chatprovider.ProviderAPIKeys, resolvedModelRoute, bool, error) {
return p.resolveGenerationModelOverride(
ctx, chat, keys, modelOpts,
summaryGenerationOverrideContext, readSummaryGenerationModelOverride,
)
}
Loading
Loading