diff --git a/README.md b/README.md index 171b8e4..406420b 100644 --- a/README.md +++ b/README.md @@ -609,11 +609,14 @@ nbctl metrics query --account-id 123e4567-e89b-12d3-a456-426614174000 --query "n #### `nbctl nubi` -Starts an interactive shell session with Nudgebee AI to ask questions and get insights. +Starts an interactive shell session or executes a single query with Nudgebee AI to ask questions and get insights. -* **Usage**: `nbctl nubi [account-id]` +* **Usage**: `nbctl nubi [account-id] [flags]` * **Arguments**: * `[account-id]` (optional): The account ID to use for the Nubi session. If not provided, `nbctl` will attempt to use the `default-account-id` from your configuration. +* **Flags**: + * `-q, --query `: Execute a single query non-interactively and exit. + * `--async`: Trigger the query asynchronously without waiting for the response (used with `--query`). * **Prerequisites**: Your `account-id` and `username` must be configured (see `nbctl configure`). **Interactive Features (within the Nubi shell):** @@ -633,14 +636,22 @@ The Nubi shell supports various slash commands to manage your session and intera * `/functions`: Lists all available functions that Nubi can execute. * `/exit`: Exits the Nubi interactive shell. -**Example:** +**Examples:** +*Start interactive session:* ```bash nbctl nubi +nbctl nubi my-dev-account-id ``` +*Single query mode (synchronous):* ```bash -nbctl nubi my-dev-account-id +nbctl nubi -q "What is the health status of my services?" +``` + +*Single query mode (asynchronous):* +```bash +nbctl nubi -q "Analyze latest deployment logs" --async ``` #### `nbctl optimizations` diff --git a/cmd/nubi.go b/cmd/nubi.go index c67473c..dba9de9 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -4,6 +4,7 @@ import ( "bufio" "context" "encoding/json" + "errors" "fmt" "os" "os/signal" @@ -105,6 +106,98 @@ var nubiCmd = &cobra.Command{ os.Exit(0) }() + query, err := cmd.Flags().GetString("query") + if err != nil { + return err + } + async, err := cmd.Flags().GetBool("async") + if err != nil { + return err + } + + query = strings.TrimSpace(query) + + if async && !cmd.Flags().Changed("query") { + return fmt.Errorf("--async requires --query / -q") + } + + // Single query mode (non-interactive) + if cmd.Flags().Changed("query") { + if query == "" { + return fmt.Errorf("query cannot be empty") + } + + ctx, cancel := context.WithCancel(cmd.Context()) + defer cancel() + + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + go func() { + select { + case <-sigChan: + cancel() + case <-ctx.Done(): + } + }() + defer signal.Stop(sigChan) + + if async { + if err := s.nubiClient.TriggerInvestigation(ctx, query); err != nil { + return fmt.Errorf("failed to trigger investigation: %w", err) + } + out := format.GetFormat().GetOutput() + _, _ = fmt.Fprintln(out, "Investigation triggered asynchronously.") + _, _ = fmt.Fprintf(out, "Session ID: %s\n", s.nubiClient.SessionID) + return nil + } + + s.spinner.Start() + startTime := time.Now() + response, status, err := s.triggerAndPoll(ctx, query) + duration := time.Since(startTime) + s.spinner.Stop() + + out := format.GetFormat().GetOutput() + grayStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + + if err != nil { + if errors.Is(err, context.Canceled) { + _, _ = fmt.Fprintln(out, "Request canceled.") + return nil + } + return fmt.Errorf("error executing query: %w", err) + } + + s.lastResponse = response + + if status == "WAITING" { + _, _ = fmt.Fprintln(out, response) + _, _ = fmt.Fprintln(out, grayStyle.Render(fmt.Sprintf("\nNote: Nubi is waiting for a followup response. To continue interactively, run 'nbctl nubi' and switch to this conversation using:\n /conversation %s\nOr visit the URL below.", s.nubiClient.ConversationID))) + } else { + rendered, err := renderMarkdown(response) + if err != nil { + _, _ = fmt.Fprintf(out, "Error rendering markdown: %v\n", err) + _, _ = fmt.Fprintln(out, response) + } else { + borderStyle := lipgloss.NewStyle().BorderStyle(lipgloss.RoundedBorder()).Padding(0, 1) + _, _ = fmt.Fprintln(out, borderStyle.Render(rendered)) + } + } + + metrics, err := s.nubiClient.GetUsageMetrics(ctx) + if err == nil && metrics != "" { + _, _ = fmt.Fprintln(out, grayStyle.Render(metrics)) + } + + _, _ = fmt.Fprintln(out, grayStyle.Render(fmt.Sprintf("Response time: %s", duration))) + + endpoint := strings.TrimSuffix(s.nubiClient.Endpoint, "/") + conversationURL := fmt.Sprintf("For more details: %s/ask-nudgebee?accountId=%s&conversation_id=%s", endpoint, s.nubiClient.AccountID, s.nubiClient.ConversationID) + _, _ = fmt.Fprintln(out, grayStyle.Render(conversationURL)) + + return nil + } + printNubiArt() // Welcome message styling @@ -600,5 +693,7 @@ func saveHistory(file string, history []string) error { } func init() { + nubiCmd.Flags().StringP("query", "q", "", "Execute a single query non-interactively and exit") + nubiCmd.Flags().Bool("async", false, "Trigger query asynchronously without waiting for response (use with --query)") rootCmd.AddCommand(nubiCmd) } diff --git a/cmd/nubi_create.go b/cmd/nubi_create.go deleted file mode 100644 index ad0ec55..0000000 --- a/cmd/nubi_create.go +++ /dev/null @@ -1,22 +0,0 @@ -package cmd - -import ( - "fmt" - - "github.com/spf13/cobra" -) - -var nubiCreateCmd = &cobra.Command{ - Use: "create", - Short: "Start a new Nubi conversation", - RunE: func(cmd *cobra.Command, args []string) error { - if _, err := fmt.Fprintln(cmd.OutOrStdout(), "nubi create called"); err != nil { - _ = err - } - return nil - }, -} - -func init() { - nubiCmd.AddCommand(nubiCreateCmd) -} diff --git a/cmd/nubi_delete.go b/cmd/nubi_delete.go deleted file mode 100644 index e7e0ce4..0000000 --- a/cmd/nubi_delete.go +++ /dev/null @@ -1,22 +0,0 @@ -package cmd - -import ( - "fmt" - - "github.com/spf13/cobra" -) - -var nubiDeleteCmd = &cobra.Command{ - Use: "delete", - Short: "Delete a Nubi conversation", - RunE: func(cmd *cobra.Command, args []string) error { - if _, err := fmt.Fprintln(cmd.OutOrStdout(), "nubi delete called"); err != nil { - _ = err - } - return nil - }, -} - -func init() { - nubiCmd.AddCommand(nubiDeleteCmd) -} diff --git a/cmd/nubi_suggest.go b/cmd/nubi_suggest.go deleted file mode 100644 index 177677a..0000000 --- a/cmd/nubi_suggest.go +++ /dev/null @@ -1,22 +0,0 @@ -package cmd - -import ( - "fmt" - - "github.com/spf13/cobra" -) - -var suggestCmd = &cobra.Command{ - Use: "suggest", - Short: "Get suggestions for a Nubi conversation", - RunE: func(cmd *cobra.Command, args []string) error { - if _, err := fmt.Fprintln(cmd.OutOrStdout(), "nubi suggest called"); err != nil { - _ = err - } - return nil - }, -} - -func init() { - nubiCmd.AddCommand(suggestCmd) -} diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go new file mode 100644 index 0000000..323a02e --- /dev/null +++ b/cmd/nubi_test.go @@ -0,0 +1,127 @@ +package cmd + +import ( + "encoding/json" + "net/http" + "testing" + + "github.com/nudgebee/nbctl/pkg/testutil" + "github.com/spf13/viper" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func resetNubiFlags() { + if f := nubiCmd.Flags().Lookup("query"); f != nil { + _ = f.Value.Set("") + f.Changed = false + } + if f := nubiCmd.Flags().Lookup("async"); f != nil { + _ = f.Value.Set("false") + f.Changed = false + } + viper.Set("username", "") +} + +func TestNubiCmd_AsyncQuery(t *testing.T) { + resetNubiFlags() + viper.Set("username", "test-user") + t.Cleanup(resetNubiFlags) + + mockResponse := map[string]interface{}{ + "ai_execute_investigation": map[string]interface{}{ + "data": map[string]interface{}{ + "response": "ok", + }, + }, + } + + output, err := testutil.RunWithSimpleGraphQL(mockResponse, nubiCmd, []string{"nubi", "test-account-id", "-q", "hello", "--async"}) + require.NoError(t, err) + + assert.Contains(t, output, "Investigation triggered asynchronously.") + assert.Contains(t, output, "Session ID:") +} + +func TestNubiCmd_AsyncWithoutQuery(t *testing.T) { + resetNubiFlags() + viper.Set("username", "test-user") + t.Cleanup(resetNubiFlags) + + _, err := testutil.RunWithSimpleGraphQL(nil, nubiCmd, []string{"nubi", "test-account-id", "--async"}) + require.Error(t, err) + assert.Contains(t, err.Error(), "--async requires --query / -q") +} + +func TestNubiCmd_EmptyQuery(t *testing.T) { + resetNubiFlags() + viper.Set("username", "test-user") + t.Cleanup(resetNubiFlags) + + _, err := testutil.RunWithSimpleGraphQL(nil, nubiCmd, []string{"nubi", "test-account-id", "-q", " "}) + require.Error(t, err) + assert.Contains(t, err.Error(), "query cannot be empty") +} + +func TestNubiCmd_SyncQuery(t *testing.T) { + resetNubiFlags() + viper.Set("username", "test-user") + t.Cleanup(resetNubiFlags) + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/api/auth/token": + _ = json.NewEncoder(w).Encode(map[string]any{"token": "fake-token", "expiry": 3600}) + case "/api/graphql": + resp := map[string]interface{}{ + "data": map[string]interface{}{ + "ai_execute_investigation": map[string]interface{}{ + "data": map[string]interface{}{ + "response": "started", + }, + }, + "ai_get_conversation_v3": map[string]interface{}{ + "conversation": map[string]interface{}{ + "id": "conv-123", + "status": "COMPLETED", + }, + "messages": []map[string]interface{}{ + { + "id": "msg-1", + "status": "COMPLETED", + "response": "System status is healthy", + "message_type": "generation", + }, + }, + }, + "ai_get_conversation_usage_metrics": map[string]interface{}{ + "data": map[string]interface{}{ + "conversation": map[string]interface{}{ + "total_cost": 0.001, + "total_input_tokens": 50, + "total_output_tokens": 100, + }, + }, + }, + }, + } + _ = json.NewEncoder(w).Encode(resp) + default: + http.NotFound(w, r) + } + }) + + defaults := map[string]any{ + "api-key": "dummy", + "username": "dummy-user", + "account-id": "dummy-account", + } + output, err := testutil.RunWithMockServer(handler, defaults, nubiCmd, []string{"nubi", "test-account-id", "-q", "system status"}) + require.NoError(t, err) + + assert.Contains(t, output, "System status") + assert.Contains(t, output, "healthy") + assert.Contains(t, output, "Cost: $0.001000") + assert.Contains(t, output, "Response time:") +} diff --git a/pkg/nubi/nubi.go b/pkg/nubi/nubi.go index f5debde..d57cae8 100644 --- a/pkg/nubi/nubi.go +++ b/pkg/nubi/nubi.go @@ -416,8 +416,12 @@ func (c *NubiClient) GetUsageMetrics(ctx context.Context) (string, error) { } `) + convID := c.ConversationID + if convID == "" { + convID = c.SessionID + } req.Var("accountId", c.AccountID) - req.Var("conversationId", c.SessionID) + req.Var("conversationId", convID) var respData struct { AiGetConversationUsageMetrics struct {