git.iko.soy/crush
on main help sign in
Merge pull request #1743 from charmbracelet/new-session
a8f5b39Christian Rocha
internal/agent/agent.go+57−27
compare with main
unifiedsplit
@@ -16,6 +16,7 @@ import (
16 "fmt"
16 "fmt"
17 "log/slog"
17 "log/slog"
18 "os"
18 "os"
19+ "regexp"
19 "strconv"
20 "strconv"
20 "strings"
21 "strings"
21 "sync"
22 "sync"
@@ -39,12 +40,17 @@ import (
39 "github.com/charmbracelet/crush/internal/stringext"
40 "github.com/charmbracelet/crush/internal/stringext"
40)
41)
41
42
43+const defaultSessionName = "Untitled Session"
44+
42//go:embed templates/title.md
45//go:embed templates/title.md
43var titlePrompt []byte
46var titlePrompt []byte
44
47
45//go:embed templates/summary.md
48//go:embed templates/summary.md
46var summaryPrompt []byte
49var summaryPrompt []byte
47
50
51+// Used to remove <think> tags from generated titles.
52+var thinkTagRegex = regexp.MustCompile(`<think>.*?</think>`)
53+
48type SessionAgentCall struct {
54type SessionAgentCall struct {
49 SessionID string
55 SessionID string
50 Prompt string
56 Prompt string
@@ -725,49 +731,68 @@ func (a *sessionAgent) getSessionMessages(ctx context.Context, session session.S
725 return msgs, nil
731 return msgs, nil
726}
732}
727
733
734+// generateTitle generates a session titled based on the initial prompt.
728−func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prompt string) {
735+func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, userPrompt string) {
729− if prompt == "" {
736+ if userPrompt == "" {
730 return
737 return
731 }
738 }
732
739
733− var maxOutput int64 = 40
740+ var maxOutputTokens int64 = 40
734 if a.smallModel.CatwalkCfg.CanReason {
741 if a.smallModel.CatwalkCfg.CanReason {
735− maxOutput = a.smallModel.CatwalkCfg.DefaultMaxTokens
742+ maxOutputTokens = a.smallModel.CatwalkCfg.DefaultMaxTokens
736 }
743 }
737
744
745+ newAgent := func(m fantasy.LanguageModel, p []byte, tok int64) fantasy.Agent {
738− agent := fantasy.NewAgent(a.smallModel.Model,
746+ return fantasy.NewAgent(m,
739− fantasy.WithSystemPrompt(string(titlePrompt)+"\n /no_think"),
747+ fantasy.WithSystemPrompt(string(p)+"\n /no_think"),
740− fantasy.WithMaxOutputTokens(maxOutput),
748+ fantasy.WithMaxOutputTokens(tok),
741− )
749+ )
750+ }
742
751
743− resp, err := agent.Stream(ctx, fantasy.AgentStreamCall{
752+ streamCall := fantasy.AgentStreamCall{
744− 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),
745− 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) {
746− prepared.Messages = options.Messages
755+ prepared.Messages = opts.Messages
747 if a.systemPromptPrefix != "" {
756 if a.systemPromptPrefix != "" {
757+ prepared.Messages = append([]fantasy.Message{
748− prepared.Messages = append([]fantasy.Message{fantasy.NewSystemMessage(a.systemPromptPrefix)}, prepared.Messages...)
758+ fantasy.NewSystemMessage(a.systemPromptPrefix),
759+ }, prepared.Messages...)
749 }
760 }
750− return callContext, prepared, nil
761+ return callCtx, prepared, nil
751 },
762 },
752− })
753− if err != nil {
754− slog.Error("error generating title", "err", err)
755− return
756 }
763 }
757
764
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+ }
759
785
760− title = strings.ReplaceAll(title, "\n", " ")
786+ title := strings.ReplaceAll(resp.Response.Content.Text(), "\n", " ")
787+ slog.Info("generated title", "title", title)
761
788
762 // Remove thinking tags if present.
789 // Remove thinking tags if present.
763− if idx := strings.Index(title, "</think>"); idx > 0 {
790+ title = thinkTagRegex.ReplaceAllString(title, "")
764− title = title[idx+len("</think>"):]
765− }
766
791
767 title = strings.TrimSpace(title)
792 title = strings.TrimSpace(title)
768 if title == "" {
793 if title == "" {
769− slog.Warn("failed to generate title", "warn", "empty title")
794+ slog.Warn("empty title; using fallback")
770− return
795+ title = defaultSessionName
771 }
796 }
772
797
773 // Calculate usage and cost.
798 // Calculate usage and cost.
@@ -783,7 +808,11 @@ func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prom
783 }
808 }
784 }
809 }
785
810
811+ if model == nil {
812+ slog.Error("no model available for cost calculation")
813+ return
814+ }
786− modelConfig := a.smallModel.CatwalkCfg
815+ modelConfig := model.CatwalkCfg
787 cost := modelConfig.CostPer1MInCached/1e6*float64(resp.TotalUsage.CacheCreationTokens) +
816 cost := modelConfig.CostPer1MInCached/1e6*float64(resp.TotalUsage.CacheCreationTokens) +
788 modelConfig.CostPer1MOutCached/1e6*float64(resp.TotalUsage.CacheReadTokens) +
817 modelConfig.CostPer1MOutCached/1e6*float64(resp.TotalUsage.CacheReadTokens) +
789 modelConfig.CostPer1MIn/1e6*float64(resp.TotalUsage.InputTokens) +
818 modelConfig.CostPer1MIn/1e6*float64(resp.TotalUsage.InputTokens) +
@@ -805,7 +834,7 @@ func (a *sessionAgent) generateTitle(ctx context.Context, sessionID string, prom
805 // concurrent session updates.
834 // concurrent session updates.
806 saveErr := a.sessions.UpdateTitleAndUsage(ctx, sessionID, title, promptTokens, completionTokens, cost)
835 saveErr := a.sessions.UpdateTitleAndUsage(ctx, sessionID, title, promptTokens, completionTokens, cost)
807 if saveErr != nil {
836 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)
809 return
838 return
810 }
839 }
811}
840}
@@ -947,6 +976,7 @@ func (a *sessionAgent) promptPrefix() string {
947 return a.systemPromptPrefix
976 return a.systemPromptPrefix
948}
977}
949
978
979+// XXX: this should be generalized to cover other subscription plans, like Copilot.
950func (a *sessionAgent) isClaudeCode() bool {
980func (a *sessionAgent) isClaudeCode() bool {
951 cfg := config.Get()
981 cfg := config.Get()
952 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 }