feat: add AI query feature support #79
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,184 @@ | ||
| package commands | ||
|
|
||
| import ( | ||
| "context" | ||
| "fmt" | ||
| "os" | ||
| "runtime" | ||
| "strings" | ||
| "time" | ||
|
|
||
| "github.com/PromptPal/go-sdk/promptpal" | ||
| "github.com/briandowns/spinner" | ||
| "github.com/gookit/color" | ||
| "github.com/sirupsen/logrus" | ||
| "github.com/urfave/cli/v2" | ||
| ) | ||
|
|
||
| var QueryCommand *cli.Command = &cli.Command{ | ||
| Name: "query", | ||
| Aliases: []string{"q"}, | ||
| Usage: "Query AI for command suggestions", | ||
| Action: commandQuery, | ||
| Description: `Query AI for command suggestions based on your prompt. | ||
|
|
||
| Examples: | ||
| shelltime query "get the top 5 memory-using processes" | ||
| shelltime q "find all files modified in the last 24 hours" | ||
| shelltime q "show disk usage for current directory"`, | ||
| } | ||
|
|
||
| func commandQuery(c *cli.Context) error { | ||
| ctx, span := commandTracer.Start(c.Context, "query") | ||
| defer span.End() | ||
|
|
||
| // Get the query from command arguments | ||
| args := c.Args().Slice() | ||
| if len(args) == 0 { | ||
| color.Red.Println("❌ Please provide a query") | ||
| return fmt.Errorf("query is required") | ||
| } | ||
|
|
||
| query := strings.Join(args, " ") | ||
|
|
||
| // Get system context | ||
| systemContext, err := getSystemContext() | ||
| if err != nil { | ||
| logrus.Warnf("Failed to get system context: %v", err) | ||
| systemContext = "Unknown system" | ||
| } | ||
|
|
||
| // Prepare the full prompt with context | ||
| fullPrompt := fmt.Sprintf(`You are a helpful assistant that suggests shell commands based on user queries. | ||
|
|
||
| System Context: | ||
| %s | ||
|
|
||
| User Query: %s | ||
|
|
||
| Please provide ONLY the shell command (no explanations, no markdown, no additional text) that would accomplish what the user is asking for. The command should be suitable for the detected operating system.`, systemContext, query) | ||
|
|
||
| s := spinner.New(spinner.CharSets[35], 200*time.Millisecond) | ||
| s.Start() | ||
| defer s.Stop() | ||
|
|
||
| // Query the AI | ||
| response, err := queryAI(ctx, fullPrompt) | ||
| if err != nil { | ||
| s.Stop() | ||
| color.Red.Printf("❌ Failed to query AI: %v\n", err) | ||
| return err | ||
| } | ||
|
|
||
| s.Stop() | ||
|
|
||
| // Display the response | ||
| color.Green.Printf("💡 Suggested command:\n") | ||
| color.Cyan.Printf("%s\n", strings.TrimSpace(response)) | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| func getSystemContext() (string, error) { | ||
| // Get current working directory | ||
| pwd, err := os.Getwd() | ||
| if err != nil { | ||
| pwd = "unknown" | ||
| } | ||
|
|
||
| // Get OS information | ||
| osInfo := runtime.GOOS | ||
|
|
||
| // Get architecture | ||
| arch := runtime.GOARCH | ||
|
|
||
| // Get hostname | ||
| hostname, err := os.Hostname() | ||
| if err != nil { | ||
| hostname = "unknown" | ||
| } | ||
|
|
||
| // Try to get some recent commands (this would be nice to have but not critical) | ||
| // For now, we'll just provide basic system info | ||
|
|
||
| context := fmt.Sprintf(`Operating System: %s | ||
| Architecture: %s | ||
| Hostname: %s | ||
| Current Working Directory: %s | ||
| Current User: %s`, osInfo, arch, hostname, pwd, os.Getenv("USER")) | ||
|
|
||
| return context, nil | ||
| } | ||
|
|
||
| func queryAI(ctx context.Context, prompt string) (string, error) { | ||
| // Create a mock configuration as requested | ||
| endpoint := "https://api.promptpal.net" // Mock URL - this would normally be configured | ||
| token := "mock-api-token" // Mock token for demonstration | ||
|
Comment on lines
+115
to
+116
Contributor
There was a problem hiding this comment. Hardcoding API endpoint and tokens is a significant security risk and makes the application difficult to configure for different environments. While this is for a demonstration, it's a good practice to avoid committing credentials. Please consider using environment variables or the existing configuration service to manage these values. |
||
|
|
||
| // Create client | ||
| oneMinute := 1 * time.Minute | ||
| promptpalClient := promptpal.NewPromptPalClient(endpoint, token, promptpal.PromptPalClientOptions{ | ||
| Timeout: &oneMinute, | ||
| }) | ||
|
|
||
| // Use a simple prompt ID for the demo - in a real scenario this would be configured | ||
| promptID := "shell-command-assistant" | ||
|
Contributor
There was a problem hiding this comment. This |
||
|
|
||
| // Variables to pass to the prompt | ||
| variables := map[string]interface{}{ | ||
| "query": prompt, | ||
| } | ||
|
|
||
| // Execute stream API as requested | ||
| var result strings.Builder | ||
| response, err := promptpalClient.ExecuteStream(ctx, promptID, variables, nil, func(data *promptpal.APIRunPromptResponse) error { | ||
| result.WriteString(data.ResponseMessage) | ||
| return nil | ||
| }) | ||
|
|
||
| if err != nil { | ||
| // For demonstration purposes, return a mock response when the API fails | ||
| // This allows the command to work even without a real PromptPal setup | ||
| return getMockResponse(prompt), nil | ||
| } | ||
|
Comment on lines
+139
to
+143
Contributor
There was a problem hiding this comment. Swallowing the error from if err != nil {
return "", err
} |
||
|
|
||
| // Return the full response | ||
| if response != nil && response.ResponseMessage != "" { | ||
| return response.ResponseMessage, nil | ||
| } | ||
|
|
||
| return result.String(), nil | ||
| } | ||
|
|
||
| // getMockResponse provides a mock AI response for demonstration purposes | ||
| func getMockResponse(query string) string { | ||
| // Simple pattern matching for common queries | ||
| lowerQuery := strings.ToLower(query) | ||
|
|
||
| if strings.Contains(lowerQuery, "memory") || strings.Contains(lowerQuery, "top") || strings.Contains(lowerQuery, "processes") { | ||
| return "ps -eo pmem,comm | sort -k 1 -r | head -5" | ||
| } | ||
|
|
||
| if strings.Contains(lowerQuery, "disk") || strings.Contains(lowerQuery, "usage") || strings.Contains(lowerQuery, "space") { | ||
| return "df -h" | ||
| } | ||
|
|
||
| if strings.Contains(lowerQuery, "files") && strings.Contains(lowerQuery, "modified") { | ||
| return "find . -type f -mtime -1" | ||
| } | ||
|
|
||
| if strings.Contains(lowerQuery, "running") && strings.Contains(lowerQuery, "processes") { | ||
| return "ps aux" | ||
| } | ||
|
|
||
| if strings.Contains(lowerQuery, "network") || strings.Contains(lowerQuery, "port") { | ||
| return "netstat -tuln" | ||
| } | ||
|
|
||
| if strings.Contains(lowerQuery, "cpu") || strings.Contains(lowerQuery, "load") { | ||
| return "top -bn1 | head -20" | ||
| } | ||
|
|
||
| // Default response | ||
| return "echo 'Unable to determine the appropriate command. Please try a more specific query.'" | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The call to
s.Stop()here is redundant. Thedefer s.Stop()on line 63 ensures the spinner is stopped when the function returns, both on success and on error. You can safely remove this line.