@@ -34,19 +34,31 @@ func commandQuery(c *cli.Context) error {
3434
3535 // Check if AI service is initialized
3636 if aiService == nil {
37- color .Red .Println ("❌ AI service is not configured" )
37+ color .Red .Println ("AI service is not configured" )
3838 return fmt .Errorf ("AI service is not available" )
3939 }
4040
4141 // Get the query from command arguments
4242 args := c .Args ().Slice ()
4343 if len (args ) == 0 {
44- color .Red .Println ("❌ Please provide a query" )
44+ color .Red .Println ("Please provide a query" )
4545 return fmt .Errorf ("query is required" )
4646 }
4747
4848 query := strings .Join (args , " " )
4949
50+ // Read config to get endpoint/token
51+ cfg , err := configService .ReadConfigFile (ctx )
52+ if err != nil {
53+ color .Red .Printf ("Failed to read config: %v\n " , err )
54+ return fmt .Errorf ("failed to read config: %w" , err )
55+ }
56+
57+ endpoint := model.Endpoint {
58+ APIEndpoint : cfg .APIEndpoint ,
59+ Token : cfg .Token ,
60+ }
61+
5062 // Get system context
5163 systemContext , err := getSystemContext (query )
5264 if err != nil {
@@ -59,34 +71,40 @@ func commandQuery(c *cli.Context) error {
5971 BaseColor : stloader.RGB {R : 100 , G : 180 , B : 255 },
6072 })
6173 l .Start ()
62- defer l .Stop ()
6374
64- // skip userId for now
65- userId := ""
75+ var result strings. Builder
76+ firstToken := true
6677
67- // Query the AI
68- newCommand , err := aiService .QueryCommand (ctx , systemContext , userId )
69- if err != nil {
78+ // Stream the AI response
79+ err = aiService .QueryCommandStream (ctx , systemContext , endpoint , func (token string ) {
80+ if firstToken {
81+ l .Stop ()
82+ color .Green .Printf ("Suggested command:\n " )
83+ firstToken = false
84+ }
85+ fmt .Print (token )
86+ result .WriteString (token )
87+ })
88+
89+ if firstToken {
90+ // No tokens received, stop loader
7091 l .Stop ()
71- color .Red .Printf ("❌ Failed to query AI: %v\n " , err )
92+ }
93+
94+ if err != nil {
95+ if ! firstToken {
96+ fmt .Println ()
97+ }
98+ color .Red .Printf ("Failed to query AI: %v\n " , err )
7299 return err
73100 }
74101
75- l .Stop ()
102+ // Print newline after streaming
103+ fmt .Println ()
76104
77- // Trim the command
78- newCommand = strings .TrimSpace (newCommand )
105+ newCommand := strings .TrimSpace (result .String ())
79106
80107 // Check auto-run configuration
81- cfg , err := configService .ReadConfigFile (ctx )
82- if err != nil {
83- slog .Warn ("Failed to read config for auto-run check" , slog .Any ("err" , err ))
84- // If can't read config, just display the command
85- displayCommand (newCommand )
86- return nil
87- }
88-
89- // Check if AI auto-run is configured
90108 if cfg .AI != nil && (cfg .AI .Agent .View || cfg .AI .Agent .Edit || cfg .AI .Agent .Delete ) {
91109 // Classify the command
92110 actionType := model .ClassifyCommand (newCommand )
@@ -105,9 +123,7 @@ func commandQuery(c *cli.Context) error {
105123 if canAutoRun {
106124 // For delete commands, add an extra confirmation
107125 if actionType == model .ActionDelete {
108- color .Green .Printf ("💡 Suggested command:\n " )
109- color .Cyan .Printf ("%s\n \n " , newCommand )
110- color .Yellow .Printf ("⚠️ This is a DELETE command. Are you sure you want to run it? (y/N): " )
126+ color .Yellow .Printf ("This is a DELETE command. Are you sure you want to run it? (y/N): " )
111127
112128 var response string
113129 fmt .Scanln (& response )
@@ -116,26 +132,20 @@ func commandQuery(c *cli.Context) error {
116132 return nil
117133 }
118134 } else {
119- // Display the command and auto-run it
120- color .Green .Printf ("💡 Auto-running command:\n " )
121- color .Cyan .Printf ("%s\n \n " , newCommand )
135+ color .Green .Printf ("Auto-running command...\n " )
122136 }
123137
124138 // Execute the command
125139 return executeCommand (ctx , newCommand )
126140 } else {
127- // Display command with info about why it's not auto-running
128- displayCommand (newCommand )
129141 if shouldShowTips (cfg ) && actionType != model .ActionOther {
130- color .Yellow .Printf ("\n 💡 Tip : This is a %s command. Enable 'ai.agent.%s' in your config to auto-run it.\n " ,
142+ color .Yellow .Printf ("\n Tip : This is a %s command. Enable 'ai.agent.%s' in your config to auto-run it.\n " ,
131143 actionType , actionType )
132144 }
133145 }
134146 } else {
135- // No auto-run configured, display the command and tip
136- displayCommand (newCommand )
137147 if shouldShowTips (cfg ) {
138- color .Yellow .Printf ("\n 💡 Tip : You can enable AI auto-run in your config file:\n " )
148+ color .Yellow .Printf ("\n Tip : You can enable AI auto-run in your config file:\n " )
139149 color .Yellow .Printf (" [ai.agent]\n " )
140150 color .Yellow .Printf (" view = true # Auto-run view commands\n " )
141151 color .Yellow .Printf (" edit = true # Auto-run edit commands\n " )
@@ -146,11 +156,6 @@ func commandQuery(c *cli.Context) error {
146156 return nil
147157}
148158
149- func displayCommand (command string ) {
150- color .Green .Printf ("💡 Suggested command:\n " )
151- color .Cyan .Printf ("%s\n " , command )
152- }
153-
154159func shouldShowTips (cfg model.ShellTimeConfig ) bool {
155160 // If ShowTips is not set (nil), default to true
156161 if cfg .AI == nil || cfg .AI .ShowTips == nil {
@@ -176,14 +181,14 @@ func executeCommand(ctx context.Context, command string) error {
176181
177182 // Run the command
178183 if err := cmd .Run (); err != nil {
179- color .Red .Printf ("\n ❌ Command failed: %v\n " , err )
184+ color .Red .Printf ("\n Command failed: %v\n " , err )
180185 return err
181186 }
182187
183188 return nil
184189}
185190
186- func getSystemContext (query string ) (model.PPPromptGuessNextPromptVariables , error ) {
191+ func getSystemContext (query string ) (model.CommandSuggestVariables , error ) {
187192 // Get shell information
188193 shell := os .Getenv ("SHELL" )
189194 if shell == "" {
@@ -198,7 +203,7 @@ func getSystemContext(query string) (model.PPPromptGuessNextPromptVariables, err
198203 // Get OS information
199204 osInfo := runtime .GOOS
200205
201- return model.PPPromptGuessNextPromptVariables {
206+ return model.CommandSuggestVariables {
202207 Shell : shell ,
203208 Os : osInfo ,
204209 Query : query ,
0 commit comments