f6a79e4 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

feat: stream content in non-interactive mode (#133)

f6a79e4Carlos Alexandro Becker1 year ago

cmd/root.go+1−17

compare with main

unifiedsplit

@@ -12,7 +12,6 @@ import (

12 "github.com/charmbracelet/crush/internal/app"

13 "github.com/charmbracelet/crush/internal/config"

14 "github.com/charmbracelet/crush/internal/db"

15− "github.com/charmbracelet/crush/internal/format"

16 "github.com/charmbracelet/crush/internal/llm/agent"

15 "github.com/charmbracelet/crush/internal/llm/agent"

17 "github.com/charmbracelet/crush/internal/log"

16 "github.com/charmbracelet/crush/internal/log"

18 "github.com/charmbracelet/crush/internal/tui"

17 "github.com/charmbracelet/crush/internal/tui"

@@ -52,14 +51,8 @@ to assist developers in writing, debugging, and understanding code directly from

52 debug, _ := cmd.Flags().GetBool("debug")

51 debug, _ := cmd.Flags().GetBool("debug")

53 cwd, _ := cmd.Flags().GetString("cwd")

52 cwd, _ := cmd.Flags().GetString("cwd")

54 prompt, _ := cmd.Flags().GetString("prompt")

53 prompt, _ := cmd.Flags().GetString("prompt")

55− outputFormat, _ := cmd.Flags().GetString("output-format")

56 quiet, _ := cmd.Flags().GetBool("quiet")

54 quiet, _ := cmd.Flags().GetBool("quiet")

58− // Validate format option

59− if !format.IsValid(outputFormat) {

60− return fmt.Errorf("invalid format option: %s\n%s", outputFormat, format.GetHelpText())

61− }

62−

63 if cwd != "" {

56 if cwd != "" {

64 err := os.Chdir(cwd)

57 err := os.Chdir(cwd)

65 if err != nil {

58 if err != nil {

@@ -109,7 +102,7 @@ to assist developers in writing, debugging, and understanding code directly from

109 // Non-interactive mode

102 // Non-interactive mode

110 if prompt != "" {

103 if prompt != "" {

111 // Run non-interactive flow using the App method

104 // Run non-interactive flow using the App method

112− return app.RunNonInteractive(ctx, prompt, outputFormat, quiet)

105+ return app.RunNonInteractive(ctx, prompt, quiet)

113 }

106 }

114

107

115 // Set up the TUI

108 // Set up the TUI

@@ -164,17 +157,8 @@ func init() {

164 rootCmd.Flags().BoolP("debug", "d", false, "Debug")

157 rootCmd.Flags().BoolP("debug", "d", false, "Debug")

165 rootCmd.Flags().StringP("prompt", "p", "", "Prompt to run in non-interactive mode")

158 rootCmd.Flags().StringP("prompt", "p", "", "Prompt to run in non-interactive mode")

166

159

167− // Add format flag with validation logic

168− rootCmd.Flags().StringP("output-format", "f", format.Text.String(),

169− "Output format for non-interactive mode (text, json)")

170−

171 // Add quiet flag to hide spinner in non-interactive mode

160 // Add quiet flag to hide spinner in non-interactive mode

172 rootCmd.Flags().BoolP("quiet", "q", false, "Hide spinner in non-interactive mode")

161 rootCmd.Flags().BoolP("quiet", "q", false, "Hide spinner in non-interactive mode")

173−

174− // Register custom validation for the format flag

175− rootCmd.RegisterFlagCompletionFunc("output-format", func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {

176− return format.SupportedFormats, cobra.ShellCompDirectiveNoFileComp

177− })

178}

162}

179

163

180func maybePrependStdin(prompt string) (string, error) {

164func maybePrependStdin(prompt string) (string, error) {

internal/app/app.go+39−25

@@ -92,7 +92,7 @@ func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) {

92}

94// RunNonInteractive handles the execution flow when a prompt is provided via CLI flag.

95−func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat string, quiet bool) error {

95+func (a *App) RunNonInteractive(ctx context.Context, prompt string, quiet bool) error {

96 slog.Info("Running in non-interactive mode")

98 // Start spinner if not in quiet mode

@@ -100,8 +100,15 @@ func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat

100 if !quiet {

101 spinner = format.NewSpinner(ctx, "Generating")

102 spinner.Start()

103− defer spinner.Stop()

104 }

103 }

104+ // Helper function to stop spinner once

105+ stopSpinner := func() {

106+ if !quiet && spinner != nil {

107+ spinner.Stop()

108+ spinner = nil

109+ }

110+ }

111+ defer stopSpinner()

105

112

106 const maxPromptLengthForTitle = 100

113 const maxPromptLengthForTitle = 100

107 titlePrefix := "Non-interactive: "

114 titlePrefix := "Non-interactive: "

@@ -128,35 +135,42 @@ func (a *App) RunNonInteractive(ctx context.Context, prompt string, outputFormat

128 return fmt.Errorf("failed to start agent processing stream: %w", err)

135 return fmt.Errorf("failed to start agent processing stream: %w", err)

129 }

136 }

130

137

131− result := <-done

138+ messageEvents := a.Messages.Subscribe(ctx)

139+ readBts := 0

132

140

133− // Stop spinner before printing output

141+ for {

134− if !quiet && spinner != nil {

142+ select {

135− spinner.Stop()

143+ case result := <-done:

144+ stopSpinner()

145+

146+ if result.Error != nil {

147+ if errors.Is(result.Error, context.Canceled) || errors.Is(result.Error, agent.ErrRequestCancelled) {

148+ slog.Info("Agent processing cancelled", "session_id", sess.ID)

149+ return nil

150+ }

151+ return fmt.Errorf("agent processing failed: %w", result.Error)

136− }

152+ }

153+

154+ part := result.Message.Content().String()[readBts:]

155+ fmt.Println(part)

137

156

138− if result.Error != nil {

139− if errors.Is(result.Error, context.Canceled) || errors.Is(result.Error, agent.ErrRequestCancelled) {

140− slog.Info("Agent processing cancelled", "session_id", sess.ID)

157+ slog.Info("Non-interactive run completed", "session_id", sess.ID)

141 return nil

158 return nil

142− }

143− return fmt.Errorf("agent processing failed: %w", result.Error)

144− }

145

159

146− // Get the text content from the response

160+ case event := <-messageEvents:

147− content := "No content available"

161+ msg := event.Payload

148− if result.Message.Content().String() != "" {

162+ if msg.SessionID == sess.ID && msg.Role == message.Assistant && len(msg.Parts) > 0 {

163+ stopSpinner()

149− content = result.Message.Content().String()

164+ part := msg.Content().String()[readBts:]

165+ fmt.Print(part)

166+ readBts += len(part)

150− }

167+ }

151

168

152− out, err := format.FormatOutput(content, outputFormat)

169+ case <-ctx.Done():

153− if err != nil {

170+ stopSpinner()

154− return err

171+ return ctx.Err()

172+ }

155 }

173 }

156−

157− fmt.Println(out)

158− slog.Info("Non-interactive run completed", "session_id", sess.ID)

159− return nil

160}

174}

161

175

162func (app *App) UpdateAgentModel() error {

176func (app *App) UpdateAgentModel() error {

internal/format/format.godeleted−91

@@ -1,91 +0,0 @@

1−package format

2−

3−import (

4− "encoding/json"

5− "fmt"

6− "strings"

7−)

8−

9−// OutputFormat represents the output format type for non-interactive mode

10−type OutputFormat string

11−

12−const (

13− // Text format outputs the AI response as plain text.

14− Text OutputFormat = "text"

15−

16− // JSON format outputs the AI response wrapped in a JSON object.

17− JSON OutputFormat = "json"

18−)

19−

20−// String returns the string representation of the OutputFormat

21−func (f OutputFormat) String() string {

22− return string(f)

23−}

24−

25−// SupportedFormats is a list of all supported output formats as strings

26−var SupportedFormats = []string{

27− string(Text),

28− string(JSON),

29−}

30−

31−// Parse converts a string to an OutputFormat

32−func Parse(s string) (OutputFormat, error) {

33− s = strings.ToLower(strings.TrimSpace(s))

34−

35− switch s {

36− case string(Text):

37− return Text, nil

38− case string(JSON):

39− return JSON, nil

40− default:

41− return "", fmt.Errorf("invalid format: %s", s)

42− }

43−}

44−

45−// IsValid checks if the provided format string is supported

46−func IsValid(s string) bool {

47− _, err := Parse(s)

48− return err == nil

49−}

50−

51−// GetHelpText returns a formatted string describing all supported formats

52−func GetHelpText() string {

53− return fmt.Sprintf(`Supported output formats:

54−- %s: Plain text output (default)

55−- %s: Output wrapped in a JSON object`,

56− Text, JSON)

57−}

58−

59−// FormatOutput formats the AI response according to the specified format

60−func FormatOutput(content string, formatStr string) (string, error) {

61− format, err := Parse(formatStr)

62− if err != nil {

63− format = Text

64− }

65−

66− switch format {

67− case JSON:

68− return formatAsJSON(content)

69− case Text:

70− fallthrough

71− default:

72− return content, nil

73− }

74−}

75−

76−// formatAsJSON wraps the content in a simple JSON object

77−func formatAsJSON(content string) (string, error) {

78− // Use the JSON package to properly escape the content

79− response := struct {

80− Response string `json:"response"`

81− }{

82− Response: content,

83− }

84−

85− jsonBytes, err := json.MarshalIndent(response, "", " ")

86− if err != nil {

87− return "", fmt.Errorf("failed to marshal output into JSON: %w", err)

88− }

89−

90− return string(jsonBytes), nil

91−}