feat: allow admins to override the chat compaction model (#27151)

This commit is contained in:
Michael Suchacz
2026-07-14 16:57:04 +02:00
committed by GitHub
parent bf57da58e7
commit a567f6a89f
30 changed files with 1454 additions and 41 deletions
+14
View File
@@ -3027,6 +3027,13 @@ func (q *querier) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (datab
return fetch(q.log, q.auth, q.db.GetChatByIDForUpdate)(ctx, id)
}
func (q *querier) GetChatCompactionModelOverride(ctx context.Context) (string, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return "", err
}
return q.db.GetChatCompactionModelOverride(ctx)
}
func (q *querier) GetChatComputerUseProvider(ctx context.Context) (string, error) {
// The computer-use provider is a deployment-wide runtime chat setting
// read by authenticated chat users and chatd. Feature and experiment
@@ -8708,6 +8715,13 @@ func (q *querier) UpsertChatAutoArchiveDays(ctx context.Context, autoArchiveDays
return q.db.UpsertChatAutoArchiveDays(ctx, autoArchiveDays)
}
func (q *querier) UpsertChatCompactionModelOverride(ctx context.Context, value string) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return err
}
return q.db.UpsertChatCompactionModelOverride(ctx, value)
}
func (q *querier) UpsertChatComputerUseProvider(ctx context.Context, provider string) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return err
+8
View File
@@ -1199,6 +1199,10 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().GetChatTitleGenerationModelOverride(gomock.Any()).Return("", nil).AnyTimes()
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead)
}))
s.Run("GetChatCompactionModelOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().GetChatCompactionModelOverride(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)
@@ -1630,6 +1634,10 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().UpsertChatTitleGenerationModelOverride(gomock.Any(), "").Return(nil).AnyTimes()
check.Args("").Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
}))
s.Run("UpsertChatCompactionModelOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().UpsertChatCompactionModelOverride(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)
+16
View File
@@ -1417,6 +1417,14 @@ func (m queryMetricsStore) GetChatByIDForUpdate(ctx context.Context, id uuid.UUI
return r0, r1
}
func (m queryMetricsStore) GetChatCompactionModelOverride(ctx context.Context) (string, error) {
start := time.Now()
r0, r1 := m.s.GetChatCompactionModelOverride(ctx)
m.queryLatencies.WithLabelValues("GetChatCompactionModelOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatCompactionModelOverride").Inc()
return r0, r1
}
func (m queryMetricsStore) GetChatComputerUseProvider(ctx context.Context) (string, error) {
start := time.Now()
r0, r1 := m.s.GetChatComputerUseProvider(ctx)
@@ -6145,6 +6153,14 @@ func (m queryMetricsStore) UpsertChatAutoArchiveDays(ctx context.Context, autoAr
return r0
}
func (m queryMetricsStore) UpsertChatCompactionModelOverride(ctx context.Context, value string) error {
start := time.Now()
r0 := m.s.UpsertChatCompactionModelOverride(ctx, value)
m.queryLatencies.WithLabelValues("UpsertChatCompactionModelOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatCompactionModelOverride").Inc()
return r0
}
func (m queryMetricsStore) UpsertChatComputerUseProvider(ctx context.Context, provider string) error {
start := time.Now()
r0 := m.s.UpsertChatComputerUseProvider(ctx, provider)
+29
View File
@@ -2608,6 +2608,21 @@ func (mr *MockStoreMockRecorder) GetChatByIDForUpdate(ctx, id any) *gomock.Call
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatByIDForUpdate", reflect.TypeOf((*MockStore)(nil).GetChatByIDForUpdate), ctx, id)
}
// GetChatCompactionModelOverride mocks base method.
func (m *MockStore) GetChatCompactionModelOverride(ctx context.Context) (string, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetChatCompactionModelOverride", ctx)
ret0, _ := ret[0].(string)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetChatCompactionModelOverride indicates an expected call of GetChatCompactionModelOverride.
func (mr *MockStoreMockRecorder) GetChatCompactionModelOverride(ctx any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatCompactionModelOverride", reflect.TypeOf((*MockStore)(nil).GetChatCompactionModelOverride), ctx)
}
// GetChatComputerUseProvider mocks base method.
func (m *MockStore) GetChatComputerUseProvider(ctx context.Context) (string, error) {
m.ctrl.T.Helper()
@@ -11514,6 +11529,20 @@ func (mr *MockStoreMockRecorder) UpsertChatAutoArchiveDays(ctx, autoArchiveDays
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatAutoArchiveDays", reflect.TypeOf((*MockStore)(nil).UpsertChatAutoArchiveDays), ctx, autoArchiveDays)
}
// UpsertChatCompactionModelOverride mocks base method.
func (m *MockStore) UpsertChatCompactionModelOverride(ctx context.Context, value string) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpsertChatCompactionModelOverride", ctx, value)
ret0, _ := ret[0].(error)
return ret0
}
// UpsertChatCompactionModelOverride indicates an expected call of UpsertChatCompactionModelOverride.
func (mr *MockStoreMockRecorder) UpsertChatCompactionModelOverride(ctx, value any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatCompactionModelOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatCompactionModelOverride), ctx, value)
}
// UpsertChatComputerUseProvider mocks base method.
func (m *MockStore) UpsertChatComputerUseProvider(ctx context.Context, provider string) error {
m.ctrl.T.Helper()
+2
View File
@@ -373,6 +373,7 @@ type sqlcQuerier interface {
GetChatByID(ctx context.Context, id uuid.UUID) (Chat, error)
GetChatByIDForShare(ctx context.Context, id uuid.UUID) (Chat, error)
GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (Chat, error)
GetChatCompactionModelOverride(ctx context.Context) (string, error)
GetChatComputerUseProvider(ctx context.Context) (string, error)
// Per-root-chat cost breakdown for a single user within a date range.
// Groups by root_chat_id so forked chats roll up under their root.
@@ -1531,6 +1532,7 @@ type sqlcQuerier interface {
// to JSON before invoking this query.
UpsertChatAdvisorConfig(ctx context.Context, value string) error
UpsertChatAutoArchiveDays(ctx context.Context, autoArchiveDays int32) error
UpsertChatCompactionModelOverride(ctx context.Context, value string) error
UpsertChatComputerUseProvider(ctx context.Context, provider string) error
// UpsertChatDebugLoggingAllowUsers updates the runtime admin setting that
// allows users to opt into chat debug logging.
+22
View File
@@ -24099,6 +24099,18 @@ func (q *sqlQuerier) GetChatAutoArchiveDays(ctx context.Context, defaultAutoArch
return auto_archive_days, err
}
const getChatCompactionModelOverride = `-- name: GetChatCompactionModelOverride :one
SELECT
COALESCE((SELECT value FROM site_configs WHERE key = 'agents_chat_compaction_model_override'), '') :: text AS model_config_id
`
func (q *sqlQuerier) GetChatCompactionModelOverride(ctx context.Context) (string, error) {
row := q.db.QueryRowContext(ctx, getChatCompactionModelOverride)
var model_config_id string
err := row.Scan(&model_config_id)
return model_config_id, err
}
const getChatComputerUseProvider = `-- name: GetChatComputerUseProvider :one
SELECT
COALESCE((SELECT value FROM site_configs WHERE key = 'agents_computer_use_provider'), '') :: text AS provider
@@ -24540,6 +24552,16 @@ func (q *sqlQuerier) UpsertChatAutoArchiveDays(ctx context.Context, autoArchiveD
return err
}
const upsertChatCompactionModelOverride = `-- name: UpsertChatCompactionModelOverride :exec
INSERT INTO site_configs (key, value) VALUES ('agents_chat_compaction_model_override', $1)
ON CONFLICT (key) DO UPDATE SET value = $1 WHERE site_configs.key = 'agents_chat_compaction_model_override'
`
func (q *sqlQuerier) UpsertChatCompactionModelOverride(ctx context.Context, value string) error {
_, err := q.db.ExecContext(ctx, upsertChatCompactionModelOverride, value)
return err
}
const upsertChatComputerUseProvider = `-- name: UpsertChatComputerUseProvider :exec
INSERT INTO site_configs (key, value) VALUES ('agents_computer_use_provider', $1)
ON CONFLICT (key) DO UPDATE SET value = $1 WHERE site_configs.key = 'agents_computer_use_provider'
+8
View File
@@ -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: GetChatCompactionModelOverride :one
SELECT
COALESCE((SELECT value FROM site_configs WHERE key = 'agents_chat_compaction_model_override'), '') :: text AS model_config_id;
-- name: UpsertChatCompactionModelOverride :exec
INSERT INTO site_configs (key, value) VALUES ('agents_chat_compaction_model_override', $1)
ON CONFLICT (key) DO UPDATE SET value = $1 WHERE site_configs.key = 'agents_chat_compaction_model_override';
-- name: GetChatDesktopEnabled :one
SELECT
COALESCE((SELECT value = 'true' FROM site_configs WHERE key = 'agents_desktop_enabled'), false) :: boolean AS enable_desktop;