diff --git a/cmd/nubi.go b/cmd/nubi.go index dba9de9..d896110 100644 --- a/cmd/nubi.go +++ b/cmd/nubi.go @@ -145,23 +145,42 @@ var nubiCmd = &cobra.Command{ if err := s.nubiClient.TriggerInvestigation(ctx, query); err != nil { return fmt.Errorf("failed to trigger investigation: %w", err) } + if format.GetFormat().Get() == "json" { + format.GetFormat().Print(map[string]interface{}{ + "message": "Investigation triggered asynchronously.", + "session_id": s.nubiClient.SessionID, + "account_id": s.nubiClient.AccountID, + "query": query, + }) + return nil + } 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() + if format.GetFormat().Get() != "json" { + s.spinner.Start() + } startTime := time.Now() response, status, err := s.triggerAndPoll(ctx, query) duration := time.Since(startTime) - s.spinner.Stop() + if s.spinner.Active() { + s.spinner.Stop() + } out := format.GetFormat().GetOutput() - grayStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) if err != nil { if errors.Is(err, context.Canceled) { + if format.GetFormat().Get() == "json" { + format.GetFormat().Print(map[string]interface{}{ + "error": "Request canceled.", + "status": "CANCELED", + }) + return nil + } _, _ = fmt.Fprintln(out, "Request canceled.") return nil } @@ -170,6 +189,42 @@ var nubiCmd = &cobra.Command{ s.lastResponse = response + endpoint := strings.TrimSuffix(s.nubiClient.Endpoint, "/") + conversationURL := fmt.Sprintf("%s/ask-nudgebee?accountId=%s&conversation_id=%s", endpoint, s.nubiClient.AccountID, s.nubiClient.ConversationID) + + if format.GetFormat().Get() == "json" { + metrics, _ := s.nubiClient.GetUsageMetrics(ctx) + details, _ := s.nubiClient.GetConversationDetails(ctx) + + var respObj any + trimmedResp := strings.TrimSpace(response) + if err := json.Unmarshal([]byte(trimmedResp), &respObj); err != nil { + respObj = response + } + + result := map[string]interface{}{ + "account_id": s.nubiClient.AccountID, + "conversation_id": s.nubiClient.ConversationID, + "session_id": s.nubiClient.SessionID, + "query": query, + "response": respObj, + "status": status, + "duration": duration.String(), + "duration_ms": duration.Milliseconds(), + "url": conversationURL, + } + if details != nil { + result["details"] = details + } + if metrics != "" { + result["metrics"] = metrics + } + format.GetFormat().Print(result) + return nil + } + + grayStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("240")) + 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))) @@ -191,9 +246,8 @@ var nubiCmd = &cobra.Command{ _, _ = 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)) + conversationURLText := fmt.Sprintf("For more details: %s", conversationURL) + _, _ = fmt.Fprintln(out, grayStyle.Render(conversationURLText)) return nil } diff --git a/cmd/nubi_test.go b/cmd/nubi_test.go index 323a02e..7dd738a 100644 --- a/cmd/nubi_test.go +++ b/cmd/nubi_test.go @@ -125,3 +125,72 @@ func TestNubiCmd_SyncQuery(t *testing.T) { assert.Contains(t, output, "Cost: $0.001000") assert.Contains(t, output, "Response time:") } + +func TestNubiCmd_SyncQuery_JSON(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", "--format", "json"}) + require.NoError(t, err) + + var result map[string]interface{} + err = json.Unmarshal([]byte(output), &result) + require.NoError(t, err) + + assert.Equal(t, "test-account-id", result["account_id"]) + assert.Equal(t, "conv-123", result["conversation_id"]) + assert.Equal(t, "system status", result["query"]) + assert.Equal(t, "System status is healthy", result["response"]) + assert.Equal(t, "COMPLETED", result["status"]) + assert.NotEmpty(t, result["url"]) +} diff --git a/cmd/root.go b/cmd/root.go index 9e08415..d7399d0 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -109,7 +109,10 @@ func init() { // Add a persistent flag for the output format. var formatVar string - rootCmd.PersistentFlags().StringVar(&formatVar, "format", "text", "Output format (json)") + rootCmd.PersistentFlags().StringVarP(&formatVar, "format", "o", "text", "Output format (json)") + var outputVar string + rootCmd.PersistentFlags().StringVar(&outputVar, "output", "", "Output format alias (json)") + _ = rootCmd.PersistentFlags().MarkHidden("output") rootCmd.PersistentFlags().String("profile", "", "Use a specific profile from your config file") _ = viper.BindPFlag("profile", rootCmd.PersistentFlags().Lookup("profile")) @@ -121,6 +124,9 @@ func init() { // Set the format from the flag. formatValue, _ := cmd.Flags().GetString("format") + if outputVal, _ := cmd.Flags().GetString("output"); outputVal != "" { + formatValue = outputVal + } format.GetFormat().Set(formatValue) // Initialize the logger. diff --git a/pkg/nubi/nubi.go b/pkg/nubi/nubi.go index d57cae8..4f700bc 100644 --- a/pkg/nubi/nubi.go +++ b/pkg/nubi/nubi.go @@ -2,6 +2,7 @@ package nubi import ( "context" + "encoding/json" "fmt" "os" "strings" @@ -361,6 +362,152 @@ func (c *NubiClient) GetConversation(ctx context.Context) (string, string, strin return finalResponse, conv.Conversation.Status, statusText, followupMessageConfig, waitingMessageID, waitingAgentID, nil } +type ConversationDetails struct { + Conversation struct { + ID string `json:"id"` + Status string `json:"status"` + } `json:"conversation"` + Messages []map[string]any `json:"messages,omitempty"` + Agents []map[string]any `json:"agents,omitempty"` + ToolCalls []map[string]any `json:"tool_calls,omitempty"` +} + +func (c *NubiClient) GetConversationDetails(ctx context.Context) (*ConversationDetails, error) { + req := client.NewRequest(` + query GetLlmConversationDetails($accountId: String!, $sessionId: String!) { + ai_get_conversation_v3(request: {account_id: $accountId, session_id: $sessionId}) { + conversation { + id + status + } + messages { + id + status + response + message_type + message_config + parent_agent_id + } + agents { + id + message_id + agent_name + status + response + } + tool_calls { + agent_id + tool_name + parameters + } + } + } + `) + + req.Var("accountId", c.AccountID) + req.Var("sessionId", c.SessionID) + + var respData struct { + AiGetConversationV3 struct { + Conversation struct { + ID string `json:"id"` + Status string `json:"status"` + } `json:"conversation"` + Messages []struct { + ID string `json:"id"` + Status string `json:"status"` + Response string `json:"response"` + MessageType string `json:"message_type"` + MessageConfig string `json:"message_config"` + ParentAgentID string `json:"parent_agent_id"` + } `json:"messages"` + Agents []struct { + ID string `json:"id"` + MessageID string `json:"message_id"` + AgentName string `json:"agent_name"` + Status string `json:"status"` + Response string `json:"response"` + } `json:"agents"` + ToolCalls []struct { + AgentID string `json:"agent_id"` + ToolName string `json:"tool_name"` + Parameters string `json:"parameters"` + } `json:"tool_calls"` + } `json:"ai_get_conversation_v3"` + } + + if err := c.Client.Run(ctx, req, &respData); err != nil { + return nil, err + } + + conv := respData.AiGetConversationV3 + details := &ConversationDetails{ + Conversation: conv.Conversation, + } + + for _, m := range conv.Messages { + msgMap := map[string]any{ + "id": m.ID, + "status": m.Status, + "message_type": m.MessageType, + "parent_agent_id": m.ParentAgentID, + } + if m.Response != "" { + var respAny any + if err := json.Unmarshal([]byte(m.Response), &respAny); err == nil { + msgMap["response"] = respAny + } else { + msgMap["response"] = m.Response + } + } + if m.MessageConfig != "" { + var cfgAny any + if err := json.Unmarshal([]byte(m.MessageConfig), &cfgAny); err == nil { + msgMap["message_config"] = cfgAny + } else { + msgMap["message_config"] = m.MessageConfig + } + } + details.Messages = append(details.Messages, msgMap) + } + + for _, a := range conv.Agents { + agentMap := map[string]any{ + "id": a.ID, + "message_id": a.MessageID, + "agent_name": a.AgentName, + "status": a.Status, + } + if a.Response != "" { + var respAny any + if err := json.Unmarshal([]byte(a.Response), &respAny); err == nil { + agentMap["response"] = respAny + } else { + agentMap["response"] = a.Response + } + } + details.Agents = append(details.Agents, agentMap) + } + + for _, t := range conv.ToolCalls { + toolMap := map[string]any{ + "agent_id": t.AgentID, + "tool_name": t.ToolName, + } + if t.Parameters != "" { + var paramsAny any + if err := json.Unmarshal([]byte(t.Parameters), ¶msAny); err == nil { + toolMap["parameters"] = paramsAny + } else { + toolMap["parameters"] = t.Parameters + } + } + details.ToolCalls = append(details.ToolCalls, toolMap) + } + + return details, nil +} + func (c *NubiClient) SendFollowupResponse(ctx context.Context, query, agentID, messageID string) error { req := client.NewRequest(` mutation AiFollowupResponse($accountId: String!, $query: String!, $conversationId: String!, $agentId: String!, $messageId: String!) {