git.iko.soy/crush
on main help sign in
fix(sessions): generate title with large model if small model fails
7315fceChristian Rocha
internal/agent/agent.go+52−24
compare with main
unifiedsplit
@@ -40,6 +40,8 @@ import (
40 "github.com/charmbracelet/crush/internal/stringext"
40 "github.com/charmbracelet/crush/internal/stringext"
41)
41)
42
42
43+const defaultSessionName = "Untitled Session"
44+
43//go:embed templates/title.md
45//go:embed templates/title.md
44var titlePrompt []byte
46var titlePrompt []byte
45
47
@@ -729,47 +731,68 @@ func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.S
729 return msgs, nil
731 return msgs, nil
730}
732}
731
733
734+// generateTitle generates a session titled based on the initial prompt.
732−func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prompt string) {
735+func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, userPrompt string) {
733− if prompt == "" {
736+ if userPrompt == "" {
734 return
737 return
735 }
738 }
736
739
737− var maxOutput int64 = 40
740+ var maxOutputTokens int64 = 40
738 if a.smallModel.CatwalkCfg.CanReason {
741 if a.smallModel.CatwalkCfg.CanReason {
739− maxOutput = a.smallModel.CatwalkCfg.DefaultMaxTokens
742+ maxOutputTokens = a.smallModel.CatwalkCfg.DefaultMaxTokens
740 }
743 }
741
744
745+ newAgent := func(m fantasy.LanguageModel, p []byte, tok int64) fantasy.Agent {
742− agent := fantasy.NewAgent(a.smallModel.Model,
746+ return fantasy.NewAgent(m,
743− fantasy.WithSystemPrompt(string(titlePrompt)+"\n /no_think"),
747+ fantasy.WithSystemPrompt(string(p)+"\n /no_think"),
744− fantasy.WithMaxOutputTokens(maxOutput),
748+ fantasy.WithMaxOutputTokens(tok),
745− )
749+ )
750+ }
746
751
747− resp, err := agent.Stream(ctx, fantasy.AgentStreamCall{
752+ streamCall := fantasy.AgentStreamCall{
748− Prompt: fmt.Sprintf("Generate a concise title for the following content:\n\n%s\n <think>\n\n</think>", prompt),
753+ Prompt: fmt.Sprintf("Generate a concise title for the following content:\n\n%s\n <think>\n\n</think>", userPrompt),
749− PrepareStep: func(callContext context.Context, options fantasy.PrepareStepFunctionOptions) (_ context.Context, prepared fantasy.PrepareStepResult, err error) {
754+ PrepareStep: func(callCtx context.Context, opts fantasy.PrepareStepFunctionOptions) (_ context.Context, prepared fantasy.PrepareStepResult, err error) {
750− prepared.Messages = options.Messages
755+ prepared.Messages = opts.Messages
751 if a.systemPromptPrefix != "" {
756 if a.systemPromptPrefix != "" {
757+ prepared.Messages = append([]fantasy.Message{
752− prepared.Messages = append([]fantasy.Message{fantasy.NewSystemMessage(a.systemPromptPrefix)}, prepared.Messages...)
758+ fantasy.NewSystemMessage(a.systemPromptPrefix),
759+ }, prepared.Messages...)
753 }
760 }
754− return callContext, prepared, nil
761+ return callCtx, prepared, nil
755 },
762 },
756− })
757− if err != nil {
758− slog.Error("error generating title", "err", err)
759− return
760 }
763 }
761
764
762− title := resp.Response.Content.Text()
765+ // Use the small model to generate the title.
766+ model := &a.smallModel
767+ agent := newAgent(model.Model, titlePrompt, maxOutputTokens)
768+ resp, err := agent.Stream(ctx, streamCall)
769+ if err == nil {
770+ // We successfully generated a title with the small model.
771+ slog.Info("generated title with small model")
772+ } else {
773+ // It didn't work. Let's try with the big model.
774+ slog.Error("error generating title with small model; trying big model", "err", err)
775+ model = &a.largeModel
776+ agent = newAgent(model.Model, titlePrompt, maxOutputTokens)
777+ resp, err = agent.Stream(ctx, streamCall)
778+ if err == nil {
779+ slog.Info("generated title with large model")
780+ } else {
781+ // Welp, the large model didn't work either.
782+ slog.Error("error generating title with large model", "err", err)
783+ }
784+ }
763
785
764− title = strings.ReplaceAll(title, "\n", " ")
786+ title := strings.ReplaceAll(resp.Response.Content.Text(), "\n", " ")
787+ slog.Info("generated title", "title", title)
765
788
766 // Remove thinking tags if present.
789 // Remove thinking tags if present.
767 title = thinkTagRegex.ReplaceAllString(title, "")
790 title = thinkTagRegex.ReplaceAllString(title, "")
768
791
769 title = strings.TrimSpace(title)
792 title = strings.TrimSpace(title)
770 if title == "" {
793 if title == "" {
771− slog.Warn("failed to generate title", "warn", "empty title")
794+ slog.Warn("empty title; using fallback")
772− return
795+ title = defaultSessionName
773 }
796 }
774
797
775 // Calculate usage and cost.
798 // Calculate usage and cost.
@@ -785,7 +808,11 @@ func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prom
785 }
808 }
786 }
809 }
787
810
811+ if model == nil {
812+ slog.Error("no model available for cost calculation")
813+ return
814+ }
788− modelConfig := a.smallModel.CatwalkCfg
815+ modelConfig := model.CatwalkCfg
789 cost := modelConfig.CostPer1MInCached/1e6*float64(resp.TotalUsage.CacheCreationTokens) +
816 cost := modelConfig.CostPer1MInCached/1e6*float64(resp.TotalUsage.CacheCreationTokens) +
790 modelConfig.CostPer1MOutCached/1e6*float64(resp.TotalUsage.CacheReadTokens) +
817 modelConfig.CostPer1MOutCached/1e6*float64(resp.TotalUsage.CacheReadTokens) +
791 modelConfig.CostPer1MIn/1e6*float64(resp.TotalUsage.InputTokens) +
818 modelConfig.CostPer1MIn/1e6*float64(resp.TotalUsage.InputTokens) +
@@ -807,7 +834,7 @@ func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prom
807 // concurrent session updates.
834 // concurrent session updates.
808 saveErr := a.sessions.UpdateTitleAndUsage(ctx, sessionID, title, promptTokens, completionTokens, cost)
835 saveErr := a.sessions.UpdateTitleAndUsage(ctx, sessionID, title, promptTokens, completionTokens, cost)
809 if saveErr != nil {
836 if saveErr != nil {
810− slog.Error("failed to save session title & usage", "error", saveErr)
837+ slog.Error("failed to save session title and usage", "error", saveErr)
811 return
838 return
812 }
839 }
813}
840}
@@ -949,6 +976,7 @@ func (a *sessionAgent) promptPrefix() string {
949 return a.systemPromptPrefix
976 return a.systemPromptPrefix
950}
977}
951
978
979+// XXX: this should be generalized to cover other subscription plans, like Copilot.
952func (a *sessionAgent) isClaudeCode() bool {
980func (a *sessionAgent) isClaudeCode() bool {
953 cfg := config.Get()
981 cfg := config.Get()
954 pc, ok := cfg.Providers.Get(a.largeModel.ModelCfg.Provider)
982 pc, ok := cfg.Providers.Get(a.largeModel.ModelCfg.Provider)
internal/tui/page/chat/chat.go+4−1
@@ -961,7 +961,10 @@ func (p *chatPage) sendMessage(text string, attachments []message.Attachment) te
961 session := p.session
961 session := p.session
962 var cmds []tea.Cmd
962 var cmds []tea.Cmd
963 if p.session.ID == "" {
963 if p.session.ID == "" {
964+ // XXX: The second argument here is the session name, which we leave
965+ // blank as it will be auto-generated. Ideally, we remove the need for
966+ // that argument entirely.
964− newSession, err := p.app.Sessions.Create(context.Background(), "New Session")
967+ newSession, err := p.app.Sessions.Create(context.Background(), "")
965 if err != nil {
968 if err != nil {
966 return util.ReportError(err)
969 return util.ReportError(err)
967 }
970 }