@@ -16,6 +16,7 @@ import (
1616 "fmt"
1717 "log/slog"
1818 "os"
19+ "regexp"
1920 "strconv"
2021 "strings"
2122 "sync"
@@ -39,12 +40,17 @@ import (
3940 "github.com/charmbracelet/crush/internal/stringext"
4041)
4142
43+const defaultSessionName = "Untitled Session"
44+
4245//go:embed templates/title.md
4346var titlePrompt []byte
4447
4548//go:embed templates/summary.md
4649var summaryPrompt []byte
4750
51+// Used to remove <think> tags from generated titles.
52+var thinkTagRegex = regexp.MustCompile(`<think>.*?</think>`)
53+
4854type SessionAgentCall struct {
4955 SessionID string
5056 Prompt string
@@ -725,49 +731,68 @@ func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.S
725731 return msgs, nil
726732}
727733
728−func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prompt string) {
729− if prompt == "" {
734+// generateTitle generates a session titled based on the initial prompt.
735+func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, userPrompt string) {
736+ if userPrompt == "" {
730737 return
731738 }
732739
733− var maxOutput int64 = 40
740+ var maxOutputTokens int64 = 40
734741 if a.smallModel.CatwalkCfg.CanReason {
735− maxOutput = a.smallModel.CatwalkCfg.DefaultMaxTokens
742+ maxOutputTokens = a.smallModel.CatwalkCfg.DefaultMaxTokens
736743 }
737744
738− agent := fantasy.NewAgent(a.smallModel.Model,
739− fantasy.WithSystemPrompt(string(titlePrompt)+"\n /no_think"),
740− fantasy.WithMaxOutputTokens(maxOutput),
741− )
745+ newAgent := func(m fantasy.LanguageModel, p []byte, tok int64) fantasy.Agent {
746+ return fantasy.NewAgent(m,
747+ fantasy.WithSystemPrompt(string(p)+"\n /no_think"),
748+ fantasy.WithMaxOutputTokens(tok),
749+ )
750+ }
742751
743− resp, err := agent.Stream(ctx, fantasy.AgentStreamCall{
744− Prompt: fmt.Sprintf("Generate a concise title for the following content:\n\n%s\n <think>\n\n</think>", prompt),
745− PrepareStep: func(callContext context.Context, options fantasy.PrepareStepFunctionOptions) (_ context.Context, prepared fantasy.PrepareStepResult, err error) {
746− prepared.Messages = options.Messages
752+ streamCall := fantasy.AgentStreamCall{
753+ Prompt: fmt.Sprintf("Generate a concise title for the following content:\n\n%s\n <think>\n\n</think>", userPrompt),
754+ PrepareStep: func(callCtx context.Context, opts fantasy.PrepareStepFunctionOptions) (_ context.Context, prepared fantasy.PrepareStepResult, err error) {
755+ prepared.Messages = opts.Messages
747756 if a.systemPromptPrefix != "" {
748− prepared.Messages = append([]fantasy.Message{fantasy.NewSystemMessage(a.systemPromptPrefix)}, prepared.Messages...)
757+ prepared.Messages = append([]fantasy.Message{
758+ fantasy.NewSystemMessage(a.systemPromptPrefix),
759+ }, prepared.Messages...)
749760 }
750− return callContext, prepared, nil
761+ return callCtx, prepared, nil
751762 },
752− })
753− if err != nil {
754− slog.Error("error generating title", "err", err)
755− return
756763 }
757764
758− 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+ }
759785
760− title = strings.ReplaceAll(title, "\n", " ")
786+ title := strings.ReplaceAll(resp.Response.Content.Text(), "\n", " ")
787+ slog.Info("generated title", "title", title)
761788
762789 // Remove thinking tags if present.
763− if idx := strings.Index(title, "</think>"); idx > 0 {
764− title = title[idx+len("</think>"):]
765− }
790+ title = thinkTagRegex.ReplaceAllString(title, "")
766791
767792 title = strings.TrimSpace(title)
768793 if title == "" {
769− slog.Warn("failed to generate title", "warn", "empty title")
770− return
794+ slog.Warn("empty title; using fallback")
795+ title = defaultSessionName
771796 }
772797
773798 // Calculate usage and cost.
@@ -783,7 +808,11 @@ func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prom
783808 }
784809 }
785810
786− modelConfig := a.smallModel.CatwalkCfg
811+ if model == nil {
812+ slog.Error("no model available for cost calculation")
813+ return
814+ }
815+ modelConfig := model.CatwalkCfg
787816 cost := modelConfig.CostPer1MInCached/1e6*float64(resp.TotalUsage.CacheCreationTokens) +
788817 modelConfig.CostPer1MOutCached/1e6*float64(resp.TotalUsage.CacheReadTokens) +
789818 modelConfig.CostPer1MIn/1e6*float64(resp.TotalUsage.InputTokens) +
@@ -805,7 +834,7 @@ func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prom
805834 // concurrent session updates.
806835 saveErr := a.sessions.UpdateTitleAndUsage(ctx, sessionID, title, promptTokens, completionTokens, cost)
807836 if saveErr != nil {
808− slog.Error("failed to save session title & usage", "error", saveErr)
837+ slog.Error("failed to save session title and usage", "error", saveErr)
809838 return
810839 }
811840}
@@ -947,6 +976,7 @@ func (a *sessionAgent) promptPrefix() string {
947976 return a.systemPromptPrefix
948977}
949978
979+// XXX: this should be generalized to cover other subscription plans, like Copilot.
950980func (a *sessionAgent) isClaudeCode() bool {
951981 cfg := config.Get()
952982 pc, ok := cfg.Providers.Get(a.largeModel.ModelCfg.Provider)