diff --git a/.goreleaser.yaml b/.goreleaser.yaml index 5213810..ab2f39f 100644 --- a/.goreleaser.yaml +++ b/.goreleaser.yaml @@ -14,7 +14,7 @@ builds: env: - CGO_ENABLED=0 ldflags: - - -s -w -X main.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.Date}} -X main.builtBy=goreleaser -X main.uptraceDsn={{.Env.UPTRACE_DSN}} -X main.ppEndpoint={{.Env.PP_ENDPOINT}} -X main.ppToken={{.Env.PP_TOKEN}} + - -s -w -X main.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.Date}} -X main.builtBy=goreleaser -X main.uptraceDsn={{.Env.UPTRACE_DSN}} - binary: "shelltime" id: mt-mac goos: @@ -26,7 +26,7 @@ builds: env: - CGO_ENABLED=0 ldflags: - - -s -w -X main.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.Date}} -X main.builtBy=goreleaser -X main.uptraceDsn={{.Env.UPTRACE_DSN}} -X main.ppEndpoint={{.Env.PP_ENDPOINT}} -X main.ppToken={{.Env.PP_TOKEN}} + - -s -w -X main.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.Date}} -X main.builtBy=goreleaser -X main.uptraceDsn={{.Env.UPTRACE_DSN}} - binary: "shelltime-daemon" id: mt-daemon-linux goos: @@ -38,7 +38,7 @@ builds: env: - CGO_ENABLED=0 ldflags: - - -s -w -X main.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.Date}} -X main.builtBy=goreleaser -X main.uptraceDsn={{.Env.UPTRACE_DSN}} -X main.ppEndpoint={{.Env.PP_ENDPOINT}} -X main.ppToken={{.Env.PP_TOKEN}} + - -s -w -X main.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.Date}} -X main.builtBy=goreleaser -X main.uptraceDsn={{.Env.UPTRACE_DSN}} - binary: "shelltime-daemon" id: mt-daemon-mac goos: @@ -50,7 +50,7 @@ builds: env: - CGO_ENABLED=0 ldflags: - - -s -w -X main.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.Date}} -X main.builtBy=goreleaser -X main.uptraceDsn={{.Env.UPTRACE_DSN}} -X main.ppEndpoint={{.Env.PP_ENDPOINT}} -X main.ppToken={{.Env.PP_TOKEN}} + - -s -w -X main.version={{.Version}} -X main.commit={{.Commit}} -X main.date={{.Date}} -X main.builtBy=goreleaser -X main.uptraceDsn={{.Env.UPTRACE_DSN}} archives: - format: tar.gz id: mt-common diff --git a/cmd/cli/main.go b/cmd/cli/main.go index 9e25b74..19d752e 100644 --- a/cmd/cli/main.go +++ b/cmd/cli/main.go @@ -19,9 +19,6 @@ var ( commit = "none" date = "unknown" uptraceDsn = "" - - ppEndpoint = "" - ppToken = "" ) func main() { @@ -64,14 +61,9 @@ func main() { model.InjectVar(version) commands.InjectVar(version, configService) - // Initialize AI service if configured - if ppEndpoint != "" && ppToken != "" { - aiService := model.NewAIService(model.AIServiceConfig{ - Endpoint: ppEndpoint, - Token: ppToken, - Timeout: 60 * time.Second, - UserToken: cfg.Token, - }) + // Initialize AI service if user has a token configured + if cfg.Token != "" { + aiService := model.NewAIService() commands.InjectAIService(aiService) } app := cli.NewApp() diff --git a/cmd/daemon/main.go b/cmd/daemon/main.go index 3865837..abfd96a 100644 --- a/cmd/daemon/main.go +++ b/cmd/daemon/main.go @@ -22,9 +22,6 @@ var ( commit = "none" date = "unknown" uptraceDsn = "" - - ppEndpoint = "" - ppToken = "" ) func main() { diff --git a/commands/query.go b/commands/query.go index 9c0ecd2..7a8f672 100644 --- a/commands/query.go +++ b/commands/query.go @@ -34,19 +34,31 @@ func commandQuery(c *cli.Context) error { // Check if AI service is initialized if aiService == nil { - color.Red.Println("āŒ AI service is not configured") + color.Red.Println("AI service is not configured") return fmt.Errorf("AI service is not available") } // Get the query from command arguments args := c.Args().Slice() if len(args) == 0 { - color.Red.Println("āŒ Please provide a query") + color.Red.Println("Please provide a query") return fmt.Errorf("query is required") } query := strings.Join(args, " ") + // Read config to get endpoint/token + cfg, err := configService.ReadConfigFile(ctx) + if err != nil { + color.Red.Printf("Failed to read config: %v\n", err) + return fmt.Errorf("failed to read config: %w", err) + } + + endpoint := model.Endpoint{ + APIEndpoint: cfg.APIEndpoint, + Token: cfg.Token, + } + // Get system context systemContext, err := getSystemContext(query) if err != nil { @@ -59,34 +71,40 @@ func commandQuery(c *cli.Context) error { BaseColor: stloader.RGB{R: 100, G: 180, B: 255}, }) l.Start() - defer l.Stop() - // skip userId for now - userId := "" + var result strings.Builder + firstToken := true - // Query the AI - newCommand, err := aiService.QueryCommand(ctx, systemContext, userId) - if err != nil { + // Stream the AI response + err = aiService.QueryCommandStream(ctx, systemContext, endpoint, func(token string) { + if firstToken { + l.Stop() + color.Green.Printf("Suggested command:\n") + firstToken = false + } + fmt.Print(token) + result.WriteString(token) + }) + + if firstToken { + // No tokens received, stop loader l.Stop() - color.Red.Printf("āŒ Failed to query AI: %v\n", err) + } + + if err != nil { + if !firstToken { + fmt.Println() + } + color.Red.Printf("Failed to query AI: %v\n", err) return err } - l.Stop() + // Print newline after streaming + fmt.Println() - // Trim the command - newCommand = strings.TrimSpace(newCommand) + newCommand := strings.TrimSpace(result.String()) // Check auto-run configuration - cfg, err := configService.ReadConfigFile(ctx) - if err != nil { - slog.Warn("Failed to read config for auto-run check", slog.Any("err", err)) - // If can't read config, just display the command - displayCommand(newCommand) - return nil - } - - // Check if AI auto-run is configured if cfg.AI != nil && (cfg.AI.Agent.View || cfg.AI.Agent.Edit || cfg.AI.Agent.Delete) { // Classify the command actionType := model.ClassifyCommand(newCommand) @@ -105,9 +123,7 @@ func commandQuery(c *cli.Context) error { if canAutoRun { // For delete commands, add an extra confirmation if actionType == model.ActionDelete { - color.Green.Printf("šŸ’” Suggested command:\n") - color.Cyan.Printf("%s\n\n", newCommand) - color.Yellow.Printf("āš ļø This is a DELETE command. Are you sure you want to run it? (y/N): ") + color.Yellow.Printf("This is a DELETE command. Are you sure you want to run it? (y/N): ") var response string fmt.Scanln(&response) @@ -116,26 +132,20 @@ func commandQuery(c *cli.Context) error { return nil } } else { - // Display the command and auto-run it - color.Green.Printf("šŸ’” Auto-running command:\n") - color.Cyan.Printf("%s\n\n", newCommand) + color.Green.Printf("Auto-running command...\n") } // Execute the command return executeCommand(ctx, newCommand) } else { - // Display command with info about why it's not auto-running - displayCommand(newCommand) if shouldShowTips(cfg) && actionType != model.ActionOther { - color.Yellow.Printf("\nšŸ’” Tip: This is a %s command. Enable 'ai.agent.%s' in your config to auto-run it.\n", + color.Yellow.Printf("\nTip: This is a %s command. Enable 'ai.agent.%s' in your config to auto-run it.\n", actionType, actionType) } } } else { - // No auto-run configured, display the command and tip - displayCommand(newCommand) if shouldShowTips(cfg) { - color.Yellow.Printf("\nšŸ’” Tip: You can enable AI auto-run in your config file:\n") + color.Yellow.Printf("\nTip: You can enable AI auto-run in your config file:\n") color.Yellow.Printf(" [ai.agent]\n") color.Yellow.Printf(" view = true # Auto-run view commands\n") color.Yellow.Printf(" edit = true # Auto-run edit commands\n") @@ -146,11 +156,6 @@ func commandQuery(c *cli.Context) error { return nil } -func displayCommand(command string) { - color.Green.Printf("šŸ’” Suggested command:\n") - color.Cyan.Printf("%s\n", command) -} - func shouldShowTips(cfg model.ShellTimeConfig) bool { // If ShowTips is not set (nil), default to true if cfg.AI == nil || cfg.AI.ShowTips == nil { @@ -176,14 +181,14 @@ func executeCommand(ctx context.Context, command string) error { // Run the command if err := cmd.Run(); err != nil { - color.Red.Printf("\nāŒ Command failed: %v\n", err) + color.Red.Printf("\nCommand failed: %v\n", err) return err } return nil } -func getSystemContext(query string) (model.PPPromptGuessNextPromptVariables, error) { +func getSystemContext(query string) (model.CommandSuggestVariables, error) { // Get shell information shell := os.Getenv("SHELL") if shell == "" { @@ -198,7 +203,7 @@ func getSystemContext(query string) (model.PPPromptGuessNextPromptVariables, err // Get OS information osInfo := runtime.GOOS - return model.PPPromptGuessNextPromptVariables{ + return model.CommandSuggestVariables{ Shell: shell, Os: osInfo, Query: query, diff --git a/commands/query_test.go b/commands/query_test.go index 871984b..abb13e5 100644 --- a/commands/query_test.go +++ b/commands/query_test.go @@ -91,16 +91,23 @@ func (s *queryTestSuite) TestQueryCommandNoArguments() { func (s *queryTestSuite) TestQueryCommandSuccess() { // Setup mocks - expectedCommand := "ls -la" query := "list all files with details" - // Mock AI service - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return(expectedCommand, nil) - - // Mock config service - no auto-run - mockedConfig := model.ShellTimeConfig{} + // Mock config service - called first for endpoint + mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", + } s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service streaming + s.mockAI.On("QueryCommandStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + onToken := args.Get(3).(func(token string)) + onToken("ls ") + onToken("-la") + }).Return(nil) + command := []string{ "shelltime-test", "query", @@ -112,20 +119,25 @@ func (s *queryTestSuite) TestQueryCommandSuccess() { } func (s *queryTestSuite) TestQueryCommandWithMultipleArgs() { - // Setup mocks - expectedCommand := "find . -name '*.go' -type f" queryParts := []string{"find", "all", "go", "files"} fullQuery := strings.Join(queryParts, " ") - // Mock AI service - s.mockAI.On("QueryCommand", mock.Anything, mock.MatchedBy(func(sc model.PPPromptGuessNextPromptVariables) bool { - return sc.Query == fullQuery - }), "").Return(expectedCommand, nil) - // Mock config service - mockedConfig := model.ShellTimeConfig{} + mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", + } s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service streaming + s.mockAI.On("QueryCommandStream", mock.Anything, mock.MatchedBy(func(sc model.CommandSuggestVariables) bool { + return sc.Query == fullQuery + }), mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + onToken := args.Get(3).(func(token string)) + onToken("find . -name '*.go' -type f") + }).Return(nil) + command := append([]string{"shelltime-test", "query"}, queryParts...) err := s.app.Run(command) @@ -135,8 +147,16 @@ func (s *queryTestSuite) TestQueryCommandWithMultipleArgs() { func (s *queryTestSuite) TestQueryCommandAIError() { query := "complex query" + // Mock config service + mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", + } + s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service error - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return("", errors.New("AI service error")) + s.mockAI.On("QueryCommandStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Return(errors.New("AI service error")) command := []string{ "shelltime-test", @@ -150,14 +170,12 @@ func (s *queryTestSuite) TestQueryCommandAIError() { } func (s *queryTestSuite) TestQueryCommandWithAutoRunView() { - expectedCommand := "ls -la" query := "list files" - // Mock AI service - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return(expectedCommand, nil) - // Mock config with auto-run enabled for view commands mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", AI: &model.AIConfig{ Agent: model.AIAgentConfig{ View: true, @@ -168,31 +186,33 @@ func (s *queryTestSuite) TestQueryCommandWithAutoRunView() { } s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service streaming + s.mockAI.On("QueryCommandStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + onToken := args.Get(3).(func(token string)) + onToken("ls -la") + }).Return(nil) + command := []string{ "shelltime-test", "query", query, } - // Note: The actual command execution would fail in test environment - // We're testing the flow up to the execution point err := s.app.Run(command) - // The error here is expected as we can't actually execute the command in tests s.Nil(err) } func (s *queryTestSuite) TestQueryCommandWithAutoRunEdit() { - expectedCommand := "sed -i 's/foo/bar/g' /tmp/file_query_command_191.txt" - query := "replace foo with bar in file.txt" + query := "write hello to file" f, _ := os.Create("/tmp/file_query_command_191.txt") f.Close() defer os.Remove(f.Name()) - // Mock AI service - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return(expectedCommand, nil) - // Mock config with auto-run enabled for edit commands mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", AI: &model.AIConfig{ Agent: model.AIAgentConfig{ View: false, @@ -203,26 +223,30 @@ func (s *queryTestSuite) TestQueryCommandWithAutoRunEdit() { } s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service streaming - use tee which works cross-platform + s.mockAI.On("QueryCommandStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + onToken := args.Get(3).(func(token string)) + onToken("echo hello | tee /tmp/file_query_command_191.txt") + }).Return(nil) + command := []string{ "shelltime-test", "query", query, } - // The command would be auto-executed if enabled, but will fail in test err := s.app.Run(command) s.Nil(err) } func (s *queryTestSuite) TestQueryCommandWithAutoRunDeleteDisabled() { - expectedCommand := "rm -rf /tmp/test" query := "delete test directory" - // Mock AI service - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return(expectedCommand, nil) - // Mock config with auto-run disabled for delete commands mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", AI: &model.AIConfig{ Agent: model.AIAgentConfig{ View: true, @@ -233,25 +257,27 @@ func (s *queryTestSuite) TestQueryCommandWithAutoRunDeleteDisabled() { } s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service streaming + s.mockAI.On("QueryCommandStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + onToken := args.Get(3).(func(token string)) + onToken("rm -rf /tmp/test") + }).Return(nil) + command := []string{ "shelltime-test", "query", query, } - // Should not auto-run delete commands when disabled err := s.app.Run(command) assert.Nil(s.T(), err) } func (s *queryTestSuite) TestQueryCommandConfigReadError() { - expectedCommand := "echo 'test'" query := "print test" - // Mock AI service - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return(expectedCommand, nil) - - // Mock config service error - should fallback gracefully + // Mock config service error - now this should return an error from commandQuery s.mockConfig.On("ReadConfigFile", mock.Anything).Return(model.ShellTimeConfig{}, errors.New("config read error")) command := []string{ @@ -261,21 +287,27 @@ func (s *queryTestSuite) TestQueryCommandConfigReadError() { } err := s.app.Run(command) - assert.Nil(s.T(), err) + assert.NotNil(s.T(), err) + assert.Contains(s.T(), err.Error(), "failed to read config") } func (s *queryTestSuite) TestQueryCommandTrimWhitespace() { - expectedCommand := " echo 'hello' \n\t" - // trimmedCommand := "echo 'hello'" query := "print hello" - // Mock AI service returning command with whitespace - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return(expectedCommand, nil) - // Mock config service - mockedConfig := model.ShellTimeConfig{} + mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", + } s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service returning command with whitespace via streaming + s.mockAI.On("QueryCommandStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + onToken := args.Get(3).(func(token string)) + onToken(" echo 'hello' \n\t") + }).Return(nil) + command := []string{ "shelltime-test", "query", @@ -284,7 +316,6 @@ func (s *queryTestSuite) TestQueryCommandTrimWhitespace() { err := s.app.Run(command) assert.Nil(s.T(), err) - // The command should be trimmed before processing } func (s *queryTestSuite) TestGetSystemContext() { @@ -317,16 +348,22 @@ func (s *queryTestSuite) TestGetSystemContext() { } func (s *queryTestSuite) TestQueryCommandWithAlias() { - expectedCommand := "ls" query := "list" - // Mock AI service - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return(expectedCommand, nil) - // Mock config service - mockedConfig := model.ShellTimeConfig{} + mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", + } s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service streaming + s.mockAI.On("QueryCommandStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + onToken := args.Get(3).(func(token string)) + onToken("ls") + }).Return(nil) + // Test using the alias "q" instead of "query" command := []string{ "shelltime-test", @@ -352,21 +389,19 @@ func (s *queryTestSuite) TestExecuteCommand() { } func (s *queryTestSuite) TestDisplayCommand() { - // This is a simple display function, just ensure it doesn't panic + // shouldShowTips is a simple function, just ensure it doesn't panic assert.NotPanics(s.T(), func() { - displayCommand("test command") + shouldShowTips(model.ShellTimeConfig{}) }) } func (s *queryTestSuite) TestQueryCommandAutoRunOtherType() { - expectedCommand := "some-complex-command --with-flags" query := "do something complex" - // Mock AI service - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return(expectedCommand, nil) - // Mock config with auto-run enabled but command is "other" type mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", AI: &model.AIConfig{ Agent: model.AIAgentConfig{ View: true, @@ -377,6 +412,13 @@ func (s *queryTestSuite) TestQueryCommandAutoRunOtherType() { } s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service streaming + s.mockAI.On("QueryCommandStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Run(func(args mock.Arguments) { + onToken := args.Get(3).(func(token string)) + onToken("some-complex-command --with-flags") + }).Return(nil) + command := []string{ "shelltime-test", "query", @@ -391,13 +433,17 @@ func (s *queryTestSuite) TestQueryCommandAutoRunOtherType() { func (s *queryTestSuite) TestQueryCommandEmptyAIResponse() { query := "do nothing" - // Mock AI service returning empty string - s.mockAI.On("QueryCommand", mock.Anything, mock.Anything, "").Return("", nil) - // Mock config service - mockedConfig := model.ShellTimeConfig{} + mockedConfig := model.ShellTimeConfig{ + APIEndpoint: "https://api.shelltime.xyz", + Token: "test-token", + } s.mockConfig.On("ReadConfigFile", mock.Anything).Return(mockedConfig, nil) + // Mock AI service returning no tokens + s.mockAI.On("QueryCommandStream", mock.Anything, mock.Anything, mock.Anything, mock.Anything). + Return(nil) + command := []string{ "shelltime-test", "query", diff --git a/go.mod b/go.mod index 3886a46..33e81a0 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,6 @@ module github.com/malamtime/cli go 1.25 require ( - github.com/PromptPal/go-sdk v0.4.2 github.com/ThreeDotsLabs/watermill v1.5.1 github.com/go-git/go-git/v5 v5.16.4 github.com/google/uuid v1.6.0 @@ -31,9 +30,6 @@ require ( atomicgo.dev/cursor v0.2.0 // indirect atomicgo.dev/keyboard v0.2.9 // indirect atomicgo.dev/schedule v0.1.0 // indirect - dario.cat/mergo v1.0.0 // indirect - github.com/Microsoft/go-winio v0.6.2 // indirect - github.com/ProtonMail/go-crypto v1.1.6 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.1 // indirect github.com/cenkalti/backoff/v5 v5.0.3 // indirect @@ -41,23 +37,18 @@ require ( github.com/clipperhouse/displaywidth v0.6.1 // indirect github.com/clipperhouse/stringish v0.1.1 // indirect github.com/clipperhouse/uax29/v2 v2.3.0 // indirect - github.com/cloudflare/circl v1.6.1 // indirect github.com/containerd/console v1.0.5 // indirect github.com/cpuguy83/go-md2man/v2 v2.0.7 // indirect github.com/cyphar/filepath-securejoin v0.6.1 // indirect github.com/davecgh/go-spew v1.1.1 // indirect - github.com/emirpasic/gods v1.18.1 // indirect github.com/fatih/color v1.18.0 // indirect github.com/felixge/httpsnoop v1.0.4 // indirect github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 // indirect github.com/go-git/go-billy/v5 v5.7.0 // indirect github.com/go-logr/logr v1.4.3 // indirect github.com/go-logr/stdr v1.2.2 // indirect - github.com/go-resty/resty/v2 v2.17.0 // indirect - github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 // indirect github.com/grpc-ecosystem/grpc-gateway/v2 v2.27.3 // indirect github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect - github.com/kevinburke/ssh_config v1.2.0 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/lithammer/fuzzysearch v1.1.8 // indirect github.com/mailru/easyjson v0.7.7 // indirect @@ -71,10 +62,8 @@ require ( github.com/pjbgf/sha1cd v0.5.0 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/russross/blackfriday/v2 v2.1.0 // indirect - github.com/skeema/knownhosts v1.3.1 // indirect github.com/stretchr/objx v0.5.3 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect - github.com/xanzy/ssh-agent v0.3.3 // indirect github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect github.com/xrash/smetrics v0.0.0-20250705151800-55b8f293f342 // indirect go.opentelemetry.io/auto/sdk v1.2.1 // indirect @@ -89,7 +78,6 @@ require ( go.opentelemetry.io/otel/sdk v1.39.0 // indirect go.opentelemetry.io/otel/sdk/log v0.15.0 // indirect go.opentelemetry.io/otel/sdk/metric v1.39.0 // indirect - golang.org/x/crypto v0.46.0 // indirect golang.org/x/net v0.48.0 // indirect golang.org/x/sys v0.39.0 // indirect golang.org/x/term v0.38.0 // indirect diff --git a/go.sum b/go.sum index ba1d4a0..d48dac6 100644 --- a/go.sum +++ b/go.sum @@ -17,19 +17,12 @@ github.com/MarvinJWendt/testza v0.3.0/go.mod h1:eFcL4I0idjtIx8P9C6KkAuLgATNKpX4/ github.com/MarvinJWendt/testza v0.4.2/go.mod h1:mSdhXiKH8sg/gQehJ63bINcCKp7RtYewEjXsvsVUPbE= github.com/MarvinJWendt/testza v0.5.2 h1:53KDo64C1z/h/d/stCYCPY69bt/OSwjq5KpFNwi+zB4= github.com/MarvinJWendt/testza v0.5.2/go.mod h1:xu53QFE5sCdjtMCKk8YMQ2MnymimEctc4n3EjyIYvEY= -github.com/Microsoft/go-winio v0.5.2/go.mod h1:WpS1mjBmmwHBEWmogvA2mj8546UReBk4v8QkMxJ6pZY= github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= -github.com/PromptPal/go-sdk v0.4.2 h1:onABRhxodJ285s7CQaXTt7cmpPQ54ztTW5FMlXfGoM0= -github.com/PromptPal/go-sdk v0.4.2/go.mod h1:67S1GmSq08wVu7Wxi//3Ru9BqcqhKcqTGey3PUJrBk8= github.com/ProtonMail/go-crypto v1.1.6 h1:ZcV+Ropw6Qn0AX9brlQLAUXfqLBc7Bl+f/DmNxpLfdw= github.com/ProtonMail/go-crypto v1.1.6/go.mod h1:rA3QumHc/FZ8pAHreoekgiAbzpNsfQAosU5td4SnOrE= github.com/ThreeDotsLabs/watermill v1.5.1 h1:t5xMivyf9tpmU3iozPqyrCZXHvoV1XQDfihas4sV0fY= github.com/ThreeDotsLabs/watermill v1.5.1/go.mod h1:Uop10dA3VeJWsSvis9qO3vbVY892LARrKAdki6WtXS4= -github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8= -github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4= -github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio= -github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs= github.com/atomicgo/cursor v0.0.1/go.mod h1:cBON2QmmrysudxNBFthvMtN32r3jxVRIvzkUiF/RuIk= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= @@ -57,16 +50,12 @@ github.com/cyphar/filepath-securejoin v0.6.1/go.mod h1:A8hd4EnAeyujCJRrICiOWqjS1 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/elazarl/goproxy v1.7.2 h1:Y2o6urb7Eule09PjlhQRGNsqRfPmYI3KKQLFpCAV3+o= -github.com/elazarl/goproxy v1.7.2/go.mod h1:82vkLNir0ALaW14Rc399OTTjyNREgmdL2cVoIbS6XaE= github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc= github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ= github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM= github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU= github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= -github.com/gliderlabs/ssh v0.3.8 h1:a4YXD1V7xMF9g5nTkdfnja3Sxy1PVDCj1Zg4Wb8vY6c= -github.com/gliderlabs/ssh v0.3.8/go.mod h1:xYoytBv1sV0aL3CavoDuJIQNURXkkfPA/wxQ1pL1fAU= github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376 h1:+zs/tPmkDkHx3U66DAb0lQFJrpS6731Oaa12ikc+DiI= github.com/go-git/gcfg v1.5.1-0.20230307220236-3a3c6141e376/go.mod h1:an3vInlBmSxCcxctByoQdvwPiA7DTK7jaaFDBTtu0ic= github.com/go-git/go-billy/v5 v5.7.0 h1:83lBUJhGWhYp0ngzCMSgllhUSuoHP1iEWYjsPl9nwqM= @@ -80,8 +69,6 @@ github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI= github.com/go-logr/logr v1.4.3/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.17.0 h1:pW9DeXcaL4Rrym4EZ8v7L19zZiIlWPg5YXAcVmt+gN0= -github.com/go-resty/resty/v2 v2.17.0/go.mod h1:kCKZ3wWmwJaNc7S29BRtUhJwy7iqmn+2mLtQrOyQlVA= github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 h1:f+oWsMOmNPc8JmEHVZIycC7hBoQxHH9pNKQORJNozsQ= github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8/go.mod h1:wcDNUvekVysuuOpQKo3191zZyTpiI6se1N1ULghS0sw= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= @@ -170,13 +157,11 @@ github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQD github.com/sergi/go-diff v1.2.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM= github.com/sergi/go-diff v1.4.0 h1:n/SP9D5ad1fORl+llWyN+D6qoUETXNZARKjyY2/KVCw= github.com/sergi/go-diff v1.4.0/go.mod h1:A0bzQcvG0E7Rwjx0REVgAGH58e96+X0MeOfepqsbeW4= -github.com/sirupsen/logrus v1.7.0/go.mod h1:yWOB1SBYBC5VeMP7gHvWumXLIWorT60ONWic61uBYv0= github.com/skeema/knownhosts v1.3.1 h1:X2osQ+RAjK76shCbvhHHHVl3ZlgDm8apHEHFqRjnBY8= github.com/skeema/knownhosts v1.3.1/go.mod h1:r7KTdC8l4uxWRyK2TpQZ/1o5HaSzh06ePQNxPwTcfiY= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.5.3 h1:jmXUvGomnU1o3W/V5h2VEradbpJDwGrzugQQvL0POH4= github.com/stretchr/objx v0.5.3/go.mod h1:rDQraq+vQZU7Fde9LOZLr8Tax6zZvy4kuNKF+QYS+U0= -github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= @@ -234,7 +219,6 @@ go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= -golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4= golang.org/x/crypto v0.46.0 h1:cKRW/pmt1pKAfetfu+RCEvjvZkA9RimPbh7bhFjGVBU= golang.org/x/crypto v0.46.0/go.mod h1:Evb/oLKmMraqjZ2iQTwDwvCtJkczlDuTmdJXoZVzqU0= golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 h1:2dVuKD2vS7b0QIHQbpyTISPd0LeHDbnYEryqj5Q1ug8= @@ -243,7 +227,6 @@ golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= -golang.org/x/net v0.0.0-20211112202133-69e39bad7dc2/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU= @@ -252,16 +235,13 @@ golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= -golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210124154548-22da62e12c0c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210330210617-4fbd30eecc44/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= -golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20211013075003-97ac67df715c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= @@ -277,14 +257,11 @@ golang.org/x/term v0.38.0 h1:PQ5pkm/rLO6HnxFR7N2lJHOZX6Kez5Y1gDSJla6jo7Q= golang.org/x/term v0.38.0/go.mod h1:bSEAKrOT1W+VSu9TSCMtoGEOUcKxOKgl3LE5QEF/xVg= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= -golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU= golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY= -golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE= -golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= diff --git a/model/ai_service.go b/model/ai_service.go index f3b318b..3a5f294 100644 --- a/model/ai_service.go +++ b/model/ai_service.go @@ -1,65 +1,93 @@ package model import ( + "bufio" + "bytes" "context" + "encoding/json" + "fmt" + "net/http" + "strings" "time" - - "github.com/PromptPal/go-sdk/promptpal" ) type AIService interface { - QueryCommand(ctx context.Context, systemContext PPPromptGuessNextPromptVariables, userId string) (string, error) + QueryCommandStream(ctx context.Context, vars CommandSuggestVariables, endpoint Endpoint, onToken func(token string)) error } -type AIServiceConfig struct { - Endpoint string - Token string - Timeout time.Duration - UserToken string +type CommandSuggestVariables struct { + Shell string `json:"shell"` + Os string `json:"os"` + Query string `json:"query"` } -type promptPalAIService struct { - client promptpal.PromptPalClient +type sseAIService struct{} + +func NewAIService() AIService { + return &sseAIService{} } -func NewAIService(config AIServiceConfig) AIService { - if config.Timeout == 0 { - config.Timeout = 1 * time.Minute +func (s *sseAIService) QueryCommandStream( + ctx context.Context, + vars CommandSuggestVariables, + endpoint Endpoint, + onToken func(token string), +) error { + body, err := json.Marshal(vars) + if err != nil { + return fmt.Errorf("failed to marshal request: %w", err) } - applyTokenFunc := func(ctx context.Context) (promptpal.ApplyTemporaryTokenResult, error) { - // Read the config to get the user's token - return promptpal.ApplyTemporaryTokenResult{ - Token: "Bearer " + config.UserToken, - }, nil - } + apiURL := strings.TrimRight(endpoint.APIEndpoint, "/") + "/api/v1/ai/command-suggest" - clientOptions := promptpal.PromptPalClientOptions{ - Timeout: &config.Timeout, - ApplyTemporaryToken: &applyTokenFunc, + req, err := http.NewRequestWithContext(ctx, http.MethodPost, apiURL, bytes.NewReader(body)) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + req.Header.Set("Authorization", "CLI "+endpoint.Token) - client := promptpal.NewPromptPalClient(config.Endpoint, config.Token, clientOptions) + client := &http.Client{Timeout: 2 * time.Minute} + resp, err := client.Do(req) + if err != nil { + return fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() - return &promptPalAIService{ - client: client, + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("server returned status %d", resp.StatusCode) } -} -func (s promptPalAIService) QueryCommand( - ctx context.Context, - systemContext PPPromptGuessNextPromptVariables, - userId string, -) (string, error) { - response, err := s.client.Execute(ctx, string(PPPromptGuessNextPrompt), PPPromptGuessNextPromptVariables{ - Shell: systemContext.Shell, - Os: systemContext.Os, - Query: systemContext.Query, - }, &userId) + scanner := bufio.NewScanner(resp.Body) + var isError bool + for scanner.Scan() { + line := scanner.Text() - if err != nil { - return "", err + if line == "event: error" { + isError = true + continue + } + + if strings.HasPrefix(line, "data: ") { + data := line[len("data: "):] + + if isError { + return fmt.Errorf("server error: %s", data) + } + + if data == "[DONE]" { + return nil + } + + onToken(data) + isError = false + } + } + + if err := scanner.Err(); err != nil { + return fmt.Errorf("error reading stream: %w", err) } - return response.ResponseMessage, nil -} \ No newline at end of file + return nil +} diff --git a/promptpal.yml b/promptpal.yml deleted file mode 100644 index bedc1eb..0000000 --- a/promptpal.yml +++ /dev/null @@ -1,10 +0,0 @@ -input: - http: - url: 'https://pp.shelltime.xyz' - token: '@env.PROMPTPAL_API_TOKEN' -output: - schema: ./schema.g.json - go_types: - prefix: PP - package_name: model - output: ./model/pp.types.g.go \ No newline at end of file