From a84869a770a9d03e38653fc78ced057c0cae41ac Mon Sep 17 00:00:00 2001 From: "claude[bot]" <209825114+claude[bot]@users.noreply.github.com> Date: Wed, 16 Jul 2025 17:36:42 +0000 Subject: [PATCH] feat: add query command with AI integration - Add PromptPal SDK dependency for AI functionality - Create query command (alias: q) for command suggestions - Implement system context gathering (OS, hostname, pwd) - Use ExecuteStream API with mock configuration - Add pattern-matching fallback for common queries - Support beautiful text output with spinner feedback Example usage: shelltime q get the top 5 memory-using processes # Returns: ps -eo pmem,comm < /dev/null | sort -k 1 -r | head -5 Co-authored-by: Le He --- cmd/cli/main.go | 1 + commands/query.go | 184 ++++++++++++++++++++++++++++++++++++++++++++++ go.mod | 4 +- go.sum | 14 +++- 4 files changed, 200 insertions(+), 3 deletions(-) create mode 100644 commands/query.go diff --git a/cmd/cli/main.go b/cmd/cli/main.go index 6dd3f98..79324b2 100644 --- a/cmd/cli/main.go +++ b/cmd/cli/main.go @@ -89,6 +89,7 @@ func main() { commands.WebCommand, commands.AliasCommand, commands.DoctorCommand, + commands.QueryCommand, } err = app.Run(os.Args) if err != nil { diff --git a/commands/query.go b/commands/query.go new file mode 100644 index 0000000..39d3f5c --- /dev/null +++ b/commands/query.go @@ -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 + + // 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" + + // 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 + } + + // 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.'" +} \ No newline at end of file diff --git a/go.mod b/go.mod index 46357fa..b4e7a4d 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/malamtime/cli go 1.24 require ( + github.com/PromptPal/go-sdk v0.4.0 github.com/ThreeDotsLabs/watermill v1.4.1 github.com/briandowns/spinner v1.23.1 github.com/gookit/color v1.5.4 @@ -13,6 +14,7 @@ require ( github.com/pkg/errors v0.9.1 github.com/sirupsen/logrus v1.9.3 github.com/stretchr/testify v1.10.0 + github.com/ugorji/go/codec v1.3.0 github.com/uptrace/uptrace-go v1.32.0 github.com/urfave/cli/v2 v2.27.4 github.com/vmihailenco/msgpack/v5 v5.4.1 @@ -30,6 +32,7 @@ require ( github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-logr/logr v1.4.2 // indirect github.com/go-logr/stdr v1.2.2 // indirect + github.com/go-resty/resty/v2 v2.7.0 // indirect github.com/google/uuid v1.6.0 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.23.0 // indirect github.com/mattn/go-colorable v0.1.2 // indirect @@ -39,7 +42,6 @@ require ( github.com/pmezard/go-difflib v1.0.0 // indirect github.com/russross/blackfriday/v2 v2.1.0 // indirect github.com/stretchr/objx v0.5.2 // indirect - github.com/ugorji/go/codec v1.3.0 // indirect github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect github.com/xo/terminfo v0.0.0-20210125001918-ca9a967f8778 // indirect github.com/xrash/smetrics v0.0.0-20240521201337-686a1a2994c1 // indirect diff --git a/go.sum b/go.sum index d257d09..3ffde04 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +github.com/PromptPal/go-sdk v0.4.0 h1:7ynJ6qPiQEchHBRK79536w3KatXSKzlslU66tt6HJv0= +github.com/PromptPal/go-sdk v0.4.0/go.mod h1:67S1GmSq08wVu7Wxi//3Ru9BqcqhKcqTGey3PUJrBk8= github.com/ThreeDotsLabs/watermill v1.4.1 h1:gjP6yZH+otMPjV0KsV07pl9TeMm9UQV/gqiuiuG5Drs= github.com/ThreeDotsLabs/watermill v1.4.1/go.mod h1:lBnrLbxOjeMRgcJbv+UiZr8Ylz8RkJ4m6i/VN/Nk+to= github.com/briandowns/spinner v1.23.1 h1:t5fDPmScwUjozhDj4FA46p5acZWIPXYE30qW2Ptu650= @@ -18,10 +20,12 @@ github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/go-resty/resty/v2 v2.7.0 h1:me+K9p3uhSmXtrBZ4k9jcEAfJmuC8IivWHwaLZwPrFY= +github.com/go-resty/resty/v2 v2.7.0/go.mod h1:9PWDzw47qPphMRFfhsyk0NnSgvluHcljSMVIq3w7q0I= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= -github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= -github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/uuid v1.2.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= @@ -115,17 +119,23 @@ go.opentelemetry.io/otel/trace v1.33.0 h1:cCJuF7LRjUFso9LPnEAHJDB2pqzp+hbO8eu1qq go.opentelemetry.io/otel/trace v1.33.0/go.mod h1:uIcdVUZMpTAmz0tI1z04GoVSezK37CbGV4fr1f2nBck= go.opentelemetry.io/proto/otlp v1.3.1 h1:TrMUixzpM0yuc/znrFTP9MMRh8trP93mkCiDVeXrui0= go.opentelemetry.io/proto/otlp v1.3.1/go.mod h1:0X1WI4de4ZsLrrJNLAQbFeLCm3T7yBkR0XqQ7niQU+8= +golang.org/x/net v0.0.0-20211029224645-99673261e6eb/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= golang.org/x/net v0.31.0 h1:68CPQngjLL0r2AlUKiSxtQFKvzRVbnzLwMUn5SzcLHo= golang.org/x/net v0.31.0/go.mod h1:P4fl1q7dY2hnZFxEk4pPSkDHF+QqjitcnDjUQyMM+pM= golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.26.0 h1:WEQa6V3Gja/BhNxg540hBip/kkaYtRg3cxg4oXSw4AU= golang.org/x/term v0.26.0/go.mod h1:Si5m1o57C5nBNQo5z1iq+XDijt21BDBDp2bK0QI8e3E= +golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.20.0 h1:gK/Kv2otX8gz+wn7Rmb3vT96ZwuoxnQlY+HlJVj7Qug= golang.org/x/text v0.20.0/go.mod h1:D4IsuqiFMhST5bX19pQ9ikHC2GsaKyk/oF+pn3ducp4= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= google.golang.org/genproto/googleapis/api v0.0.0-20241104194629-dd2ea8efbc28 h1:M0KvPgPmDZHPlbRbaNU1APr28TvwvvdUPlSv7PUvy8g= google.golang.org/genproto/googleapis/api v0.0.0-20241104194629-dd2ea8efbc28/go.mod h1:dguCy7UOdZhTvLzDyt15+rOrawrpM4q7DD9dQ1P11P4= google.golang.org/genproto/googleapis/rpc v0.0.0-20241104194629-dd2ea8efbc28 h1:XVhgTWWV3kGQlwJHR3upFWZeTsei6Oks1apkZSeonIE=