135aceb · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

feat: add dynamic model switching with agent provider updates

135acebKujtim Hoxha1 year ago

internal/app/app.go+4

compare with main

unifiedsplit

@@ -168,3 +168,7 @@ func (app *App) Shutdown() {

168 }

169 app.CoderAgent.CancelAll()

170}

171+

172+func (app *App) UpdateAgentModel() error {

173+ return app.CoderAgent.UpdateModel()

174+}

internal/llm/agent/agent.go+98

@@ -56,6 +56,7 @@ type Service interface {

56 IsSessionBusy(sessionID string) bool

57 IsBusy() bool

58 Summarize(ctx context.Context, sessionID string) error

59+ UpdateModel() error

59}

60}

61type agent struct {

62type agent struct {

@@ -805,3 +806,100 @@ func (a *agent) CancelAll() {

805 return true

806 return true

806 })

807 })

807}

808}

809+

810+func (a *agent) UpdateModel() error {

811+ cfg := config.Get()

812+

813+ // Get current provider configuration

814+ currentProviderCfg := config.GetAgentProvider(a.agentCfg.ID)

815+ if currentProviderCfg.ID == "" {

816+ return fmt.Errorf("provider for agent %s not found in config", a.agentCfg.Name)

817+ }

818+

819+ // Check if provider has changed

820+ if string(currentProviderCfg.ID) != a.providerID {

821+ // Provider changed, need to recreate the main provider

822+ model := config.GetAgentModel(a.agentCfg.ID)

823+ if model.ID == "" {

824+ return fmt.Errorf("model not found for agent %s", a.agentCfg.Name)

825+ }

826+

827+ promptID := agentPromptMap[a.agentCfg.ID]

828+ if promptID == "" {

829+ promptID = prompt.PromptDefault

830+ }

831+

832+ opts := []provider.ProviderClientOption{

833+ provider.WithModel(a.agentCfg.Model),

834+ provider.WithSystemMessage(prompt.GetPrompt(promptID, currentProviderCfg.ID)),

835+ provider.WithMaxTokens(model.DefaultMaxTokens),

836+ }

837+

838+ newProvider, err := provider.NewProviderV2(currentProviderCfg, opts...)

839+ if err != nil {

840+ return fmt.Errorf("failed to create new provider: %w", err)

841+ }

842+

843+ // Update the provider and provider ID

844+ a.provider = newProvider

845+ a.providerID = string(currentProviderCfg.ID)

846+ }

847+

848+ // Check if small model provider has changed (affects title and summarize providers)

849+ smallModelCfg := cfg.Models.Small

850+ var smallModelProviderCfg config.ProviderConfig

851+

852+ for _, p := range cfg.Providers {

853+ if p.ID == smallModelCfg.Provider {

854+ smallModelProviderCfg = p

855+ break

856+ }

857+ }

858+

859+ if smallModelProviderCfg.ID == "" {

860+ return fmt.Errorf("provider %s not found in config", smallModelCfg.Provider)

861+ }

862+

863+ // Check if summarize provider has changed

864+ if string(smallModelProviderCfg.ID) != a.summarizeProviderID {

865+ var smallModel config.Model

866+ for _, m := range smallModelProviderCfg.Models {

867+ if m.ID == smallModelCfg.ModelID {

868+ smallModel = m

869+ break

870+ }

871+ }

872+ if smallModel.ID == "" {

873+ return fmt.Errorf("model %s not found in provider %s", smallModelCfg.ModelID, smallModelProviderCfg.ID)

874+ }

875+

876+ // Recreate title provider

877+ titleOpts := []provider.ProviderClientOption{

878+ provider.WithModel(config.SmallModel),

879+ provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),

880+ provider.WithMaxTokens(40),

881+ }

882+ newTitleProvider, err := provider.NewProviderV2(smallModelProviderCfg, titleOpts...)

883+ if err != nil {

884+ return fmt.Errorf("failed to create new title provider: %w", err)

885+ }

886+

887+ // Recreate summarize provider

888+ summarizeOpts := []provider.ProviderClientOption{

889+ provider.WithModel(config.SmallModel),

890+ provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, smallModelProviderCfg.ID)),

891+ provider.WithMaxTokens(smallModel.DefaultMaxTokens),

892+ }

893+ newSummarizeProvider, err := provider.NewProviderV2(smallModelProviderCfg, summarizeOpts...)

894+ if err != nil {

895+ return fmt.Errorf("failed to create new summarize provider: %w", err)

896+ }

897+

898+ // Update the providers and provider ID

899+ a.titleProvider = newTitleProvider

900+ a.summarizeProvider = newSummarizeProvider

901+ a.summarizeProviderID = string(smallModelProviderCfg.ID)

902+ }

903+

904+ return nil

905+}

internal/tui/tui.go+7

@@ -173,6 +173,13 @@ func (a *appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

173 // Model Switch

174 case models.ModelSelectedMsg:

175 config.UpdatePreferredModel(config.LargeModel, msg.Model)

176+

177+ // Update the agent with the new model/provider configuration

178+ if err := a.app.UpdateAgentModel(); err != nil {

179+ logging.ErrorPersist(fmt.Sprintf("Failed to update agent model: %v", err))

180+ return a, util.ReportError(fmt.Errorf("model changed to %s but failed to update agent: %v", msg.Model.ModelID, err))

181+ }

182+

176 return a, util.ReportInfo(fmt.Sprintf("Model changed to %s", msg.Model.ModelID))

183 return a, util.ReportInfo(fmt.Sprintf("Model changed to %s", msg.Model.ModelID))

177

184

178 // File Picker

185 // File Picker