git.iko.soy/crush
on main help sign in
Merge pull request #557 from charmbracelet/change-summarize-model
fac2d43Kujtim Hoxha
internal/llm/agent/agent.go+37−31
compare with main
unifiedsplit
@@ -159,11 +159,12 @@ func NewAgent(
159 if err != nil {
159 if err != nil {
160 return nil, err
160 return nil, err
161 }
161 }
162+
162 summarizeOpts := []provider.ProviderClientOption{
163 summarizeOpts := []provider.ProviderClientOption{
163− provider.WithModel(config.SelectedModelTypeSmall),
164+ provider.WithModel(config.SelectedModelTypeLarge),
164− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, smallModelProviderCfg.ID)),
165+ provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, providerCfg.ID)),
165 }
166 }
166− summarizeProvider, err := provider.NewProvider(*smallModelProviderCfg, summarizeOpts...)
167+ summarizeProvider, err := provider.NewProvider(*providerCfg, summarizeOpts...)
167 if err != nil {
168 if err != nil {
168 return nil, err
169 return nil, err
169 }
170 }
@@ -224,7 +225,7 @@ func NewAgent(
224 sessions: sessions,
225 sessions: sessions,
225 titleProvider: titleProvider,
226 titleProvider: titleProvider,
226 summarizeProvider: summarizeProvider,
227 summarizeProvider: summarizeProvider,
227− summarizeProviderID: string(smallModelProviderCfg.ID),
228+ summarizeProviderID: string(providerCfg.ID),
228 activeRequests: csync.NewMap[string, context.CancelFunc](),
229 activeRequests: csync.NewMap[string, context.CancelFunc](),
229 tools: csync.NewLazySlice(toolFn),
230 tools: csync.NewLazySlice(toolFn),
230 }, nil
231 }, nil
@@ -904,54 +905,59 @@ func (a *agent) UpdateModel() error {
904 a.providerID = string(currentProviderCfg.ID)
905 a.providerID = string(currentProviderCfg.ID)
905 }
906 }
906
907
907− // Check if small model provider has changed (affects title and summarize providers)
908+ // Check if providers have changed for title (small) and summarize (large)
908 smallModelCfg := cfg.Models[config.SelectedModelTypeSmall]
909 smallModelCfg := cfg.Models[config.SelectedModelTypeSmall]
909 var smallModelProviderCfg config.ProviderConfig
910 var smallModelProviderCfg config.ProviderConfig
910−
911 for p := range cfg.Providers.Seq() {
911 for p := range cfg.Providers.Seq() {
912 if p.ID == smallModelCfg.Provider {
912 if p.ID == smallModelCfg.Provider {
913 smallModelProviderCfg = p
913 smallModelProviderCfg = p
914 break
914 break
915 }
915 }
916 }
916 }
917−
918 if smallModelProviderCfg.ID == "" {
917 if smallModelProviderCfg.ID == "" {
919 return fmt.Errorf("provider %s not found in config", smallModelCfg.Provider)
918 return fmt.Errorf("provider %s not found in config", smallModelCfg.Provider)
920 }
919 }
921
920
922− // Check if summarize provider has changed
921+ largeModelCfg := cfg.Models[config.SelectedModelTypeLarge]
923− if string(smallModelProviderCfg.ID) != a.summarizeProviderID {
922+ var largeModelProviderCfg config.ProviderConfig
924− smallModel := cfg.GetModelByType(config.SelectedModelTypeSmall)
923+ for p := range cfg.Providers.Seq() {
925− if smallModel == nil {
924+ if p.ID == largeModelCfg.Provider {
926− return fmt.Errorf("model %s not found in provider %s", smallModelCfg.Model, smallModelProviderCfg.ID)
925+ largeModelProviderCfg = p
926+ break
927 }
927 }
928+ }
929+ if largeModelProviderCfg.ID == "" {
930+ return fmt.Errorf("provider %s not found in config", largeModelCfg.Provider)
931+ }
928
932
929− // Recreate title provider
933+ // Recreate title provider
930− titleOpts := []provider.ProviderClientOption{
934+ titleOpts := []provider.ProviderClientOption{
931− provider.WithModel(config.SelectedModelTypeSmall),
935+ provider.WithModel(config.SelectedModelTypeSmall),
932− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),
936+ provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptTitle, smallModelProviderCfg.ID)),
933− // We want the title to be short, so we limit the max tokens
934− provider.WithMaxTokens(40),
937+ provider.WithMaxTokens(40),
935− }
938+ }
936− newTitleProvider, err := provider.NewProvider(smallModelProviderCfg, titleOpts...)
939+ newTitleProvider, err := provider.NewProvider(smallModelProviderCfg, titleOpts...)
937− if err != nil {
940+ if err != nil {
938− return fmt.Errorf("failed to create new title provider: %w", err)
941+ return fmt.Errorf("failed to create new title provider: %w", err)
939− }
942+ }
943+ a.titleProvider = newTitleProvider
940
944
941− // Recreate summarize provider
945+ // Recreate summarize provider if provider changed (now large model)
946+ if string(largeModelProviderCfg.ID) != a.summarizeProviderID {
947+ largeModel := cfg.GetModelByType(config.SelectedModelTypeLarge)
948+ if largeModel == nil {
949+ return fmt.Errorf("model %s not found in provider %s", largeModelCfg.Model, largeModelProviderCfg.ID)
950+ }
942 summarizeOpts := []provider.ProviderClientOption{
951 summarizeOpts := []provider.ProviderClientOption{
943− provider.WithModel(config.SelectedModelTypeSmall),
952+ provider.WithModel(config.SelectedModelTypeLarge),
944− provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, smallModelProviderCfg.ID)),
953+ provider.WithSystemMessage(prompt.GetPrompt(prompt.PromptSummarizer, largeModelProviderCfg.ID)),
945 }
954 }
946− newSummarizeProvider, err := provider.NewProvider(smallModelProviderCfg, summarizeOpts...)
955+ newSummarizeProvider, err := provider.NewProvider(largeModelProviderCfg, summarizeOpts...)
947 if err != nil {
956 if err != nil {
948 return fmt.Errorf("failed to create new summarize provider: %w", err)
957 return fmt.Errorf("failed to create new summarize provider: %w", err)
949 }
958 }
950−
951− // Update the providers and provider ID
952− a.titleProvider = newTitleProvider
953 a.summarizeProvider = newSummarizeProvider
959 a.summarizeProvider = newSummarizeProvider
954− a.summarizeProviderID = string(smallModelProviderCfg.ID)
960+ a.summarizeProviderID = string(largeModelProviderCfg.ID)
955 }
961 }
956
962
957 return nil
963 return nil