diff --git a/ai.go b/ai.go index 4e6b09c9..e2cbadbe 100644 --- a/ai.go +++ b/ai.go @@ -8,11 +8,11 @@ import ( "encoding/base64" "encoding/hex" "encoding/json" - "unicode/utf8" "errors" "fmt" "io/ioutil" "log" + "math" "math/rand" "net/http" "net/url" @@ -24,7 +24,8 @@ import ( "strings" "sync" "time" - "math" + "unicode/utf8" + openai "github.com/sashabaranov/go-openai" uuid "github.com/satori/go.uuid" "google.golang.org/api/customsearch/v1" @@ -51,6 +52,7 @@ var model = "gpt-5.4-nano" var fallbackModel = "" var assistantId = os.Getenv("OPENAI_ASSISTANT_ID") var docsVectorStoreID = os.Getenv("OPENAI_DOCS_VS_ID") +var skipAgentWait = os.Getenv("SHUFFLE_SKIP_AGENT_WAIT") var skipAgentWait = os.Getenv("SHUFFLE_SKIP_AGENT_WAIT") var agentRunLocation = os.Getenv("SHUFFLE_AGENT_RUN_LOCATION") var assistantModel = model @@ -66,19 +68,19 @@ func init() { } reasoningEffort := os.Getenv("AI_REASONING_EFFORT") - if reasoningEffort == "minimal" || reasoningEffort == "low" || reasoningEffort == "medium" || reasoningEffort == "high" { + if reasoningEffort == "minimal" || reasoningEffort == "low" || reasoningEffort == "medium" || reasoningEffort == "high" { aiReasoningEffort = reasoningEffort } } func EstimatePromptTokens(messages []openai.ChatCompletionMessage) int64 { - totalChars := 0 - for _, msg := range messages { - totalChars += utf8.RuneCountInString(msg.Content) - totalChars += 20 - } - - return int64((totalChars + 3) / 4) + totalChars := 0 + for _, msg := range messages { + totalChars += utf8.RuneCountInString(msg.Content) + totalChars += 20 + } + + return int64((totalChars + 3) / 4) } // Provide an incident triage and response plan for the reported incident finding. Make a short list of actions to perform in the following format: [{"title": "Title of the task", "category": "triage/containment/recovery/communication/documentation", "completed": false, "createdBy": "ai-agent@shuffler.io"}]. ONLY output as JSON array and nothing more. After the list is made, add these to the metadata.extensions.custom_attributes.tasks[] in the next action. @@ -2228,12 +2230,12 @@ Do not add explanations, comments, or extra formatting. Only return valid JSON.` } func GetActionAIResponse(ctx context.Context, resp http.ResponseWriter, user User, org Org, outputFormat string, input QueryInput) ([]byte, error) { - if len(org.Id) == 0 { + if len(org.Id) == 0 { if len(input.OrgId) > 0 && user.ActiveOrg.Id == "" { user.ActiveOrg.Id = input.OrgId } - if len(user.ActiveOrg.Id) > 0 { + if len(user.ActiveOrg.Id) > 0 { newOrg, err := GetOrg(ctx, user.ActiveOrg.Id) if err != nil { log.Printf("[ERROR] Failed to load orgid '%s' in ai response check", user.ActiveOrg.Id) @@ -2253,7 +2255,7 @@ func GetActionAIResponse(ctx context.Context, resp http.ResponseWriter, user Use if project.Environment == "cloud" && !user.SupportAccess { //if org.SyncFeatures.ShuffleGPT.Active && org.SyncFeatures.ShuffleGPT.Usage < org.SyncFeatures.ShuffleGPT.Limit { - // Most should never reach this + // Most should never reach this if org.SyncFeatures.ShuffleGPT.Usage < 1000 { log.Printf("[AUDIT] Org %#v (%s) has access to the auto feature. Allowing user %s to use it", org.Name, org.Id, user.Username) org.SyncFeatures.ShuffleGPT.Usage += 1 @@ -7187,8 +7189,19 @@ func sendAITokenLimitAlert(ctx context.Context, execution WorkflowExecution, ful aiPercentage := float64(monthlyTokensUsed) / float64(tokenLimit) * 100 pctCacheKey := generateAlertCacheKey(billingOrgId, fmt.Sprintf("ai_token_pct_%d", int64(100)), admins) - if !checkAndSetAlertCache(ctx, pctCacheKey) { - log.Printf("[DEBUG] Skipping duplicate AI token alert for org %s, threshold %d%% - already sent recently", billingOrgId, int64(aiPercentage)) + if !checkAndSetAlertCache(ctx, cacheKey) { + log.Printf("[DEBUG] Skipping duplicate AI token limit alert for org %s - already sent recently (1)", billingOrgId) + return + } + + orgStats, err := GetOrgStatistics(ctx, billingOrgId) + if err != nil { + log.Printf("[ERROR] Failed to get org stats for AI token limit alert for org %s: %s", billingOrgId, err) + return + } + + if orgStats.MonthlyAIUsageAlertSent { + log.Printf("[DEBUG] Skipping duplicate AI token limit alert for org %s - already sent recently (2)", billingOrgId) return } @@ -7230,6 +7243,11 @@ func sendAITokenLimitAlert(ctx context.Context, execution WorkflowExecution, ful log.Printf("[ERROR] Failed sending AI token alert email to %v for org %s: %s", admins, billingOrgId, err) } else { log.Printf("[INFO] Sent AI token %d%% alert email to %v of org %s", int64(aiPercentage), admins, billingOrgId) + orgStats.MonthlyAIUsageAlertSent = true + errStats := SetOrgStatistics(ctx, *orgStats, billingOrgId) + if errStats != nil { + log.Printf("[ERROR] Failed to update org stats after sending AI token limit alert for org %s: %s", billingOrgId, errStats) + } } } @@ -7435,15 +7453,15 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, aiStarttime := time.Now().UnixMilli() replacedExecution, err := GetWorkflowExecution(ctx, execution.ExecutionId) - if err == nil && len(replacedExecution.Results) > 0 && (execution.Status == "EXECUTING" || execution.Status == "WAITING") { + if err == nil && len(replacedExecution.Results) > 0 && (execution.Status == "EXECUTING" || execution.Status == "WAITING") { execution = *replacedExecution } - llmResponse := []byte{} - if len(aiResponseWrapper) > 0 { - if len(aiResponseWrapper[0]) > 0 { + llmResponse := []byte{} + if len(aiResponseWrapper) > 0 { + if len(aiResponseWrapper[0]) > 0 { llmResponse = aiResponseWrapper[0] - //createNextActions = false + //createNextActions = false } } @@ -7458,8 +7476,8 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, break } - if execution.Status != "EXECUTING" && execution.Status != "WAITING" { - return startNode, errors.New("Agent run already finished") + if execution.Status != "EXECUTING" && execution.Status != "WAITING" { + return startNode, errors.New("Agent run already finished") } // Metadata = org-specific context @@ -7472,25 +7490,24 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, metadata += fmt.Sprintf("Current time: %s\n", time.Now().Format(time.RFC3339)) /* - categoryActions := GetAppCategories() - actionMetadata := "ALL Available actions sorted by category:\n" - for _, category := range categoryActions { - if category.Name == "AI" || category.Name == "Other" { - continue - } + categoryActions := GetAppCategories() + actionMetadata := "ALL Available actions sorted by category:\n" + for _, category := range categoryActions { + if category.Name == "AI" || category.Name == "Other" { + continue + } - actionMetadata += "\nCategory: " + category.Name + "\n" - for _, label := range category.ActionLabels { - actionMetadata += fmt.Sprintf("- %s\n", strings.ReplaceAll(label, "_", " ")) + actionMetadata += "\nCategory: " + category.Name + "\n" + for _, label := range category.ActionLabels { + actionMetadata += fmt.Sprintf("- %s\n", strings.ReplaceAll(label, "_", " ")) + } } - } */ if len(execution.Workflow.OrgId) == 0 && len(execution.ExecutionOrg) > 0 { execution.Workflow.OrgId = execution.ExecutionOrg } - // Validate On-Prem Configuration immediately if project.Environment != "cloud" { if os.Getenv("AI_MODEL") == "" && os.Getenv("OPENAI_MODEL") == "" { @@ -7633,11 +7650,11 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, enableQuestions = true } - if param.Name == "reasoning" { + if param.Name == "reasoning" { foundReasoning = strings.ToLower(strings.TrimSpace(param.Value)) } - if param.Name == "image" { + if param.Name == "image" { if strings.HasPrefix(param.Value, "http://") || strings.HasPrefix(param.Value, "https://") { imagesIncluded = append(imagesIncluded, param.Value) } else { @@ -7650,7 +7667,7 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, } } - if param.Name == "image_detail" { + if param.Name == "image_detail" { if param.Value == "low" { imageDetail = openai.ImageURLDetailLow } else if param.Value == "high" { @@ -7663,7 +7680,7 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, } if param.Name == "app_name" { - //if debug { + //if debug { // log.Printf("[DEBUG] Rewriting app_name to action") //} @@ -7700,6 +7717,8 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, requiredParams := []string{} optionalParams := []string{} for _, param := range sortedAppAction.Parameters { + if param.Name == "body" && len(param.Example) > 0 { + if len(param.Example) > 150 { if param.Name == "url" { continue } @@ -7720,7 +7739,7 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, } if param.Required { - if param.Configuration && param.Name != "url" { + if param.Configuration && param.Name != "url" { continue } @@ -7761,7 +7780,7 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, sortedAppAction.Description = sortedAppAction.Description[:100] + "..." } descString = fmt.Sprintf(" # %s", sortedAppAction.Description) - } + } if descString == previousDesc { descString = "" @@ -7837,7 +7856,7 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, err := json.Unmarshal([]byte(result.Result), &mappedResult) if err != nil { log.Printf("[ERROR][%s] AI Agent (1): Failed unmarshalling result for action %s: %s", execution.ExecutionId, startNode.ID, err) - if debug { + if debug { log.Printf("[WARNING] FAILED AI AGENT THING: %s", result.Result) } break @@ -7860,7 +7879,7 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, if debug { log.Printf("[DEBUG][%s] Found existing WAITING decision at index %d (action=%s) - returning existing state", execution.ExecutionId, mappedDecision.I, mappedDecision.Action) } - + hasActiveDecision = true break } else if status == "RUNNING" { @@ -8065,7 +8084,7 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, } if len(foundUserId) > 0 { - foundUser, err := GetUser(ctx, foundUserId) + foundUser, err := GetUser(ctx, foundUserId) if err == nil && len(foundUser.Id) > 0 { if len(foundUser.UserGeoInfo.Country.Name) > 0 { metadata += fmt.Sprintf("Country: %s,", foundUser.UserGeoInfo.Country.Name) @@ -8099,141 +8118,141 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, // if the user doesn't want to run anything /* - decidedApps := "" - appauth, autherr := GetAllWorkflowAppAuth(ctx, org.Id) - if autherr == nil && len(appauth) > 0 { - preferredApps := []WorkflowApp{ - WorkflowApp{ - Categories: []string{"internal"}, - Name: "shuffle datastore", - }, - } - if len(org.SecurityFramework.SIEM.Name) > 0 { - preferredApps = append(preferredApps, WorkflowApp{ - Categories: []string{"siem"}, - Name: org.SecurityFramework.SIEM.Name, - }) - } - - if len(org.SecurityFramework.EDR.Name) > 0 { - //preferredApps += strings.ToLower(org.SecurityFramework.EDR.Name) + ", " - preferredApps = append(preferredApps, WorkflowApp{ - Categories: []string{"eradication"}, - Name: org.SecurityFramework.EDR.Name, - }) - } + decidedApps := "" + appauth, autherr := GetAllWorkflowAppAuth(ctx, org.Id) + if autherr == nil && len(appauth) > 0 { + preferredApps := []WorkflowApp{ + WorkflowApp{ + Categories: []string{"internal"}, + Name: "shuffle datastore", + }, + } + if len(org.SecurityFramework.SIEM.Name) > 0 { + preferredApps = append(preferredApps, WorkflowApp{ + Categories: []string{"siem"}, + Name: org.SecurityFramework.SIEM.Name, + }) + } - if len(org.SecurityFramework.Communication.Name) > 0 { - //preferredApps += strings.ToLower(org.SecurityFramework.Cases.Name) + ", " + if len(org.SecurityFramework.EDR.Name) > 0 { + //preferredApps += strings.ToLower(org.SecurityFramework.EDR.Name) + ", " + preferredApps = append(preferredApps, WorkflowApp{ + Categories: []string{"eradication"}, + Name: org.SecurityFramework.EDR.Name, + }) + } - preferredApps = append(preferredApps, WorkflowApp{ - Categories: []string{"cases"}, - Name: org.SecurityFramework.Communication.Name, - }) - } + if len(org.SecurityFramework.Communication.Name) > 0 { + //preferredApps += strings.ToLower(org.SecurityFramework.Cases.Name) + ", " - if len(org.SecurityFramework.Cases.Name) > 0 { - //preferredApps += strings.ToLower(org.SecurityFramework.Cases.Name) + ", " + preferredApps = append(preferredApps, WorkflowApp{ + Categories: []string{"cases"}, + Name: org.SecurityFramework.Communication.Name, + }) + } - preferredApps = append(preferredApps, WorkflowApp{ - Categories: []string{"cases"}, - Name: org.SecurityFramework.Cases.Name, - }) - } + if len(org.SecurityFramework.Cases.Name) > 0 { + //preferredApps += strings.ToLower(org.SecurityFramework.Cases.Name) + ", " - if len(org.SecurityFramework.Assets.Name) > 0 { - //preferredApps += strings.ToLower(org.SecurityFramework.Assets.Name) + ", " + preferredApps = append(preferredApps, WorkflowApp{ + Categories: []string{"cases"}, + Name: org.SecurityFramework.Cases.Name, + }) + } - preferredApps = append(preferredApps, WorkflowApp{ - Categories: []string{"assets"}, - Name: org.SecurityFramework.Assets.Name, - }) - } + if len(org.SecurityFramework.Assets.Name) > 0 { + //preferredApps += strings.ToLower(org.SecurityFramework.Assets.Name) + ", " - if len(org.SecurityFramework.Network.Name) > 0 { - //preferredApps += strings.ToLower(org.SecurityFramework.Network.Name) + ", " + preferredApps = append(preferredApps, WorkflowApp{ + Categories: []string{"assets"}, + Name: org.SecurityFramework.Assets.Name, + }) + } - preferredApps = append(preferredApps, WorkflowApp{ - Categories: []string{"network"}, - Name: org.SecurityFramework.Network.Name, - }) - } + if len(org.SecurityFramework.Network.Name) > 0 { + //preferredApps += strings.ToLower(org.SecurityFramework.Network.Name) + ", " - if len(org.SecurityFramework.Intel.Name) > 0 { - //preferredApps += strings.ToLower(org.SecurityFramework.Intel.Name) + ", " + preferredApps = append(preferredApps, WorkflowApp{ + Categories: []string{"network"}, + Name: org.SecurityFramework.Network.Name, + }) + } - preferredApps = append(preferredApps, WorkflowApp{ - Categories: []string{"intel"}, - Name: org.SecurityFramework.Intel.Name, - }) - } + if len(org.SecurityFramework.Intel.Name) > 0 { + //preferredApps += strings.ToLower(org.SecurityFramework.Intel.Name) + ", " - if len(org.SecurityFramework.IAM.Name) > 0 { - //preferredApps += strings.ToLower(org.SecurityFramework.IAM.Name) + ", " - preferredApps = append(preferredApps, WorkflowApp{ - Categories: []string{"iam"}, - Name: org.SecurityFramework.IAM.Name, - }) - } + preferredApps = append(preferredApps, WorkflowApp{ + Categories: []string{"intel"}, + Name: org.SecurityFramework.Intel.Name, + }) + } - for _, auth := range appauth { - // ALWAYS append valid auth - if !auth.Validation.Valid { - continue + if len(org.SecurityFramework.IAM.Name) > 0 { + //preferredApps += strings.ToLower(org.SecurityFramework.IAM.Name) + ", " + preferredApps = append(preferredApps, WorkflowApp{ + Categories: []string{"iam"}, + Name: org.SecurityFramework.IAM.Name, + }) } - if len(auth.App.Categories) > 0 { - found := false - for _, preApp := range preferredApps { - if len(preApp.Categories) == 0 { - continue + for _, auth := range appauth { + // ALWAYS append valid auth + if !auth.Validation.Valid { + continue + } + + if len(auth.App.Categories) > 0 { + found := false + for _, preApp := range preferredApps { + if len(preApp.Categories) == 0 { + continue + } + + if ArrayContains(preApp.Categories, strings.ToLower(auth.App.Categories[0])) { + found = true + break + } } - if ArrayContains(preApp.Categories, strings.ToLower(auth.App.Categories[0])) { - found = true - break + if found { + continue } } - if found { + if len(auth.App.Categories) > 0 && strings.ToUpper(auth.App.Categories[0]) == "AI" { continue } - } - if len(auth.App.Categories) > 0 && strings.ToUpper(auth.App.Categories[0]) == "AI" { - continue + preferredApps = append(preferredApps, auth.App) } - preferredApps = append(preferredApps, auth.App) - } + // FIXME: Pre-filter before this to ensure we have good + // apps ONLY. + for _, preferredApp := range preferredApps { + if len(preferredApp.Name) == 0 { + continue + } - // FIXME: Pre-filter before this to ensure we have good - // apps ONLY. - for _, preferredApp := range preferredApps { - if len(preferredApp.Name) == 0 { - continue - } + lowername := strings.ToLower(preferredApp.Name) + if strings.Contains(decidedApps, lowername) { + continue + } - lowername := strings.ToLower(preferredApp.Name) - if strings.Contains(decidedApps, lowername) { - continue + decidedApps += lowername + ", " } - decidedApps += lowername + ", " + // Let's inject http. + if !strings.Contains(decidedApps, "http") { + decidedApps += "http, " + } } - // Let's inject http. - if !strings.Contains(decidedApps, "http") { - decidedApps += "http, " + if len(decidedApps) > 0 { + // if len(allowedActionString) == 0 { + // metadata += fmt.Sprintf("\n\nALL TOOLS: %s\n\n", decidedApps) + // } + metadata += fmt.Sprintf("\n\nALL TOOLS: %s\n\n", decidedApps) } - } - - if len(decidedApps) > 0 { - // if len(allowedActionString) == 0 { - // metadata += fmt.Sprintf("\n\nALL TOOLS: %s\n\n", decidedApps) - // } - metadata += fmt.Sprintf("\n\nALL TOOLS: %s\n\n", decidedApps) - } */ } } @@ -8246,22 +8265,22 @@ func HandleAiAgentExecutionStart(execution WorkflowExecution, startNode Action, //metadata += "\n" + actionMetadata } - // Due to usually NOT wanting a question back, but pure run + // Due to usually NOT wanting a question back, but pure run enableQuestionsString := ` 2. **Explicit 'Ask' Command:** - Avoid asking questions. Have an action bias and make decisions for the user! ` - if enableQuestions { + if enableQuestions { enableQuestionsString = ` 5. **Explicit 'Ask' Command:** - **Trigger:** LOWEST PRIORITY. Does the user explicitly COMMAND you to ask them for input (e.g., "Ask me for the IP")? - **Action:** Select "ask" (Category: "standalone"). - **Field "question":** The specific questions you have. Make decisions FOR the user instead of asking. Do NOT ask questions about authentication or authorization. Do NOT ask to confirm the obvious. Assume you are allowed to use the mentioned tool. Do NOT ask unless absolutely necessary. This command should generally be avoided in favor of action bias. Have as few questions as possible, but if multiple questions are required, ask one question at a time as such: "fields": [{"key": "question", "value": "question1"}, {"key": "question", "value": "question2"}]` - // FIXME: Uncomment below and add to the enableQuestionsString. New feature for auto-generating and approving new apps. The generate API docs API supports this + // FIXME: Uncomment below and add to the enableQuestionsString. New feature for auto-generating and approving new apps. The generate API docs API supports this - // If the tool is not mentioned in USER CONTEXT and you NEED them to allow those tools, set "action": "add_tool" and "tool": "EXACT toolname" and do not ask questions. If multiple tools are required, make multiple decisions - one for each required tool. Put the entire reasoning in the "reason" field - not as fields. + // If the tool is not mentioned in USER CONTEXT and you NEED them to allow those tools, set "action": "add_tool" and "tool": "EXACT toolname" and do not ask questions. If multiple tools are required, make multiple decisions - one for each required tool. Put the entire reasoning in the "reason" field - not as fields. } systemMessage += fmt.Sprintf(`### MISSION @@ -8367,7 +8386,7 @@ data_filter: agentReasoningEffort = foundReasoning } - if skipAgentWait == "true" { + if skipAgentWait == "true" { } else if len(userMessage) == 0 { log.Printf("[ERROR][%s] AI Agent: No user message/input found for action %s", execution.ExecutionId, startNode.ID) return abortAgentExecution(ctx, execution, startNode, AgentOutput{}, "no_user_message", "No user message/input found for AI Agent start") @@ -8379,7 +8398,7 @@ data_filter: initiatedBy = "system" } - if len(llmResponse) > 0 { + if len(llmResponse) > 0 { } else if !createNextActions { if strings.TrimSpace(callerName) == "" { callerName = "unknown" @@ -8413,7 +8432,7 @@ data_filter: }, // Move towards determinism - Temperature: 0, + Temperature: 0, ReasoningEffort: agentReasoningEffort, // Reasoning control @@ -8430,10 +8449,10 @@ data_filter: }) } else { newMessage := openai.ChatCompletionMessage{ - Role: openai.ChatMessageRoleUser, + Role: openai.ChatMessageRoleUser, MultiContent: []openai.ChatMessagePart{ openai.ChatMessagePart{ - Type: openai.ChatMessagePartTypeText, + Type: openai.ChatMessagePartTypeText, Text: preparedContent, }, }, @@ -8441,9 +8460,9 @@ data_filter: for _, imageIncluded := range imagesIncluded { newMessage.MultiContent = append(newMessage.MultiContent, openai.ChatMessagePart{ - Type: openai.ChatMessagePartTypeImageURL, + Type: openai.ChatMessagePartTypeImageURL, ImageURL: &openai.ChatMessageImageURL{ - URL: imageIncluded, + URL: imageIncluded, Detail: imageDetail, }, }) @@ -8477,8 +8496,8 @@ data_filter: }) } - if len(marshalledDecisions) > 4 { - completionRequest.Messages = append(completionRequest.Messages, openai.ChatCompletionMessage { + if len(marshalledDecisions) > 4 { + completionRequest.Messages = append(completionRequest.Messages, openai.ChatCompletionMessage{ Role: openai.ChatMessageRoleUser, Content: fmt.Sprintf("HISTORY:\n%s", string(marshalledDecisions)), }) @@ -8490,7 +8509,7 @@ data_filter: Content: failureInjection, }) } - + // Let's try to make the prompt cache key sticky type ExtendedRequest struct { openai.ChatCompletionRequest @@ -8678,13 +8697,13 @@ data_filter: client.Timeout = time.Minute * 5 - // Test for whether we can ignore response wait time + // Test for whether we can ignore response wait time // This is to drastically reduce CPU use of Agent requests // 1 second = enough to read the body, which is the only major // obstacle - if skipAgentWait == "true" { - //client.Timeout = time.Second * 1 - client.Timeout = time.Millisecond * 1000 + if skipAgentWait == "true" { + //client.Timeout = time.Second * 1 + client.Timeout = time.Millisecond * 1000 fullUrl += "&skip_result_wait=true" } else { // Makes sure we wait as long as possible @@ -8715,7 +8734,7 @@ data_filter: log.Printf("[INFO][%s] Started AI Agent action %s with app '%s'. Waiting for results...", execution.ExecutionId, startNode.ID, chosenAiApp) if err != nil { - if skipAgentWait == "true" && strings.Contains(strings.ToLower(err.Error()), "timeout") { + if skipAgentWait == "true" && strings.Contains(strings.ToLower(err.Error()), "timeout") { // Question when we return here: // How do we get back to EXACTLY here when the AI is done? // Point being: we need the same data anyway. @@ -8805,7 +8824,7 @@ data_filter: continue } - if debug { + if debug { log.Printf("[DEBUG][%s] AI Agent: Found body parameter which MAY contain the right user input. LEN: %d", execution.ExecutionId, len(param.Value)) } @@ -8897,14 +8916,14 @@ data_filter: // Edgecase handling for LLM not being available etc if len(choicesString) > 0 { - if debug { + if debug { log.Printf("[ERROR][%s] AI Agent: Found choicesString (1) in AI Agent response error handling: %s", execution.ExecutionId, choicesString) } } else if len(openaiOutput.Choices) == 0 { log.Printf("[ERROR][%s] AI Agent: No choices found in AI agent response (1). Status: %d. Raw: %s", execution.ExecutionId, outputMap.Status, bodyString) - // This is specific to OpenAI, but may work for others + // This is specific to OpenAI, but may work for others newOutput := openai.ErrorResponse{} err = json.Unmarshal(bodyString, &newOutput) if err == nil && len(newOutput.Error.Message) > 0 { @@ -9013,6 +9032,7 @@ data_filter: } } + // LLM is occasionally appending freeform text like (e.g. "Summary: ...") after the closing bracket. Truncate everything past the last ']' so the JSON // LLM is occasionally appending freeform text like (e.g. "Summary: ...") after the closing bracket. Truncate everything past the last ']' so the JSON // parser doesn't dont break due to that. if lastBracket := strings.LastIndex(decisionString, "]"); lastBracket != -1 { @@ -9035,7 +9055,7 @@ data_filter: log.Printf("[ERROR][%s] AI Agent (6): Failed unmarshalling decisions in AI Agent response (2): %s. String: %s", execution.ExecutionId, err, decisionString) // Updating the OUTPUT in some way to help the user a bit. - if strings.Contains(decisionString, "conditions must be correct") { + if strings.Contains(decisionString, "conditions must be correct") { errorMessage = fmt.Sprintf("Condition failed. See decision_string for details") resultMapping.Status = "SKIPPED" } else { @@ -9145,7 +9165,7 @@ data_filter: execution.Results[resultIndex].Result = string(agentOutputMarshalled) } - // Waiting 1 + // Waiting 1 execution.Results[resultIndex].Status = "WAITING" // Update the result in cache as actions are self-corrective @@ -9251,7 +9271,7 @@ data_filter: err = CreateOrgNotification( ctx, fmt.Sprintf("Agent - approval required for '%s'", mappedDecision.Tool), - fmt.Sprintf("Approval required during agent run."), + fmt.Sprintf("Approval required during agent run."), fmt.Sprintf("/forms/%s?authorization=%s&reference_execution=%s&source_node=%s&decision_id=%s&backend_url=%s", execution.WorkflowId, execution.Authorization, execution.ExecutionId, startNode.ID, mappedDecision.RunDetails.Id, backendUrl), execution.ExecutionOrg, false, @@ -9319,7 +9339,7 @@ data_filter: log.Printf("[DEBUG][%s] AI Agent: Decision index %d is an 'ask' action. Setting approval required to true for manual review in the UI.", execution.ExecutionId, mappedDecision.I) question := mappedDecision.Reason - if len(mappedDecision.Fields) > 0 { + if len(mappedDecision.Fields) > 0 { question = mappedDecision.Fields[0].Value } @@ -9683,7 +9703,6 @@ func GenerateSingulWorkflows(resp http.ResponseWriter, request *http.Request) { if categoryAction.ActionName == "remove" || categoryAction.ActionName == "disable" || categoryAction.ActionName == "stop" { - if workflowErr == nil && workflow.OrgId == user.ActiveOrg.Id { // Delete the workflow err = DeleteKey(ctx, "workflow", workflowId, user.ActiveOrg.Id) @@ -9692,13 +9711,13 @@ func GenerateSingulWorkflows(resp http.ResponseWriter, request *http.Request) { } /* - if debug { - log.Printf("[DEBUG] DELETING KEY: %s", deleteKey) - allWorkflows, err := GetAllWorkflowsByQuery(ctx, user, 250, "") - if err == nil { - log.Printf("\n\n[DEBUG] FOUND WORKFLOWS AFTER DELETE: %d\n\n", len(allWorkflows)) + if debug { + log.Printf("[DEBUG] DELETING KEY: %s", deleteKey) + allWorkflows, err := GetAllWorkflowsByQuery(ctx, user, 250, "") + if err == nil { + log.Printf("\n\n[DEBUG] FOUND WORKFLOWS AFTER DELETE: %d\n\n", len(allWorkflows)) + } } - } */ } else { log.Printf("[INFO] No existing workflow with ID %s to remove for category '%s'", workflowId, categoryAction.Label) @@ -9909,7 +9928,7 @@ func GenerateSingulWorkflows(resp http.ResponseWriter, request *http.Request) { if len(workflow.Actions[actionIndex].LargeImage) == 0 { - if strings.Contains(strings.ToLower(action.AppName), "agent") || strings.Contains(strings.ToLower(action.AppName), "singul") || strings.Contains(strings.ToLower(action.AppName), "integration") { + if strings.Contains(strings.ToLower(action.AppName), "agent") || strings.Contains(strings.ToLower(action.AppName), "singul") || strings.Contains(strings.ToLower(action.AppName), "integration") { workflow.Actions[actionIndex].LargeImage = "/icons/workflow-page/shuffle_agent.png" } else if debug { log.Printf("[DEBUG] Missing app image for app '%s'", action.AppName) @@ -9934,12 +9953,12 @@ func GenerateSingulWorkflows(resp http.ResponseWriter, request *http.Request) { } /* - if debug { - allWorkflows, err := GetAllWorkflowsByQuery(ctx, user, 250, "") - if err == nil { - log.Printf("\n\n[DEBUG] FOUND WORKFLOWS POST CREATE: %d\n\n", len(allWorkflows)) + if debug { + allWorkflows, err := GetAllWorkflowsByQuery(ctx, user, 250, "") + if err == nil { + log.Printf("\n\n[DEBUG] FOUND WORKFLOWS POST CREATE: %d\n\n", len(allWorkflows)) + } } - } */ resp.WriteHeader(http.StatusOK) @@ -10123,7 +10142,7 @@ func RunAiQuery(ctx context.Context, info AiCallInfo, systemMessage, userMessage } } - if len(newMessages) > 5 { + if len(newMessages) > 5 { chatCompletion.Messages = newMessages } } @@ -13988,22 +14007,22 @@ func RunMCPAction(resp http.ResponseWriter, request *http.Request) { // Run the action newAction := Action{ - Name: "agent", - AppName: "AI Agent", - AppID: "shuffle_agent", - AppVersion: "1.0.0", + Name: "agent", + AppName: "AI Agent", + AppID: "shuffle_agent", + AppVersion: "1.0.0", Environment: foundEnvironment, Parameters: []WorkflowAppActionParameter{ WorkflowAppActionParameter{ - Name: "app_name", + Name: "app_name", Value: "openai", }, WorkflowAppActionParameter{ - Name: "input", + Name: "input", Value: foundRequest.Params.Input.Text, }, WorkflowAppActionParameter{ - Name: "app_name", + Name: "app_name", Value: parsedApp, }, }, @@ -14140,13 +14159,13 @@ func HandleMCPMethodInitialize(request MCPRequest, user User, app WorkflowApp) ( foundServerVersion := "0.0.1" tools := MCPInitResponse{ Jsonrpc: request.Jsonrpc, - ID: request.ID, + ID: request.ID, Result: MCPToolResult{ ProtocolVersion: "2024-11-05", - Tools: []MCPTool{}, - Capabilities: MCPCapabilities{}, + Tools: []MCPTool{}, + Capabilities: MCPCapabilities{}, ServerInfo: MCPServerInfo{ - Name: "shuffle", + Name: "shuffle", Version: foundServerVersion, }, }, @@ -14154,11 +14173,11 @@ func HandleMCPMethodInitialize(request MCPRequest, user User, app WorkflowApp) ( for cnt, action := range app.Actions { tool := MCPTool{ - Name: action.Name, + Name: action.Name, Description: action.Description, InputSchema: MCPToolInputSchema{ - Type: "object", - Required: []string{}, + Type: "object", + Required: []string{}, Properties: map[string]MCPProperty{}, }, } @@ -14183,12 +14202,12 @@ func HandleMCPMethodInitialize(request MCPRequest, user User, app WorkflowApp) ( } parsedDescription := param.Description - if strings.Contains(parsedDescription, "Generated by") { + if strings.Contains(parsedDescription, "Generated by") { parsedDescription = "" } tool.InputSchema.Properties[param.Name] = MCPProperty{ - Type: "string", + Type: "string", Description: parsedDescription, } } diff --git a/db-connector.go b/db-connector.go index d0390494..d4d42de2 100755 --- a/db-connector.go +++ b/db-connector.go @@ -160,10 +160,11 @@ func SetOrgStatistics(ctx context.Context, stats ExecutionInfo, id string) error log.Printf("[ERROR] Failed adding stats with ID %s: %s", id, putErr) if strings.Contains(fmt.Sprintf("%s", putErr), "entity is too big") { - log.Printf("[WARNING] SetOrgStatistics: entity too big for org %s – archiving to GCS and trimming", id) + log.Printf("[WARNING] SetOrgStatistics: entity too big for org %s – attempting to archive to GCS", id) if archiveErr := archiveOldStatsToGCSBucket(ctx, id, &stats); archiveErr != nil { - log.Printf("[WARNING] SetOrgStatistics: GCS archive failed for org %s: %s – trimming anyway", id, archiveErr) + log.Printf("[ERROR] SetOrgStatistics: GCS archive failed for org %s: %s – cannot trim stats without backup, returning original error", id, archiveErr) + return putErr } if len(stats.DailyStatistics) > 60 { diff --git a/stats.go b/stats.go index 8a65532c..894baaf4 100755 --- a/stats.go +++ b/stats.go @@ -846,7 +846,7 @@ func HandleGetStatistics(resp http.ResponseWriter, request *http.Request) { // Get a max of the last 365 days if len(info.DailyStatistics) > 365 { - info.DailyStatistics = info.DailyStatistics[len(info.DailyStatistics)-60:] + info.DailyStatistics = info.DailyStatistics[len(info.DailyStatistics)-365:] } } @@ -1406,6 +1406,7 @@ func handleDailyCacheUpdate(executionInfo *ExecutionInfo) *ExecutionInfo { executionInfo.MonthlyAgentOutputTokens = 0 executionInfo.LastMonthlyResetMonth = currentMonth executionInfo.LastUsageAlertThreshold = 0 + executionInfo.MonthlyAIUsageAlertSent = false // Reset all usage alerts to unsent for index := range executionInfo.UsageAlerts { diff --git a/structs.go b/structs.go index 13a86a98..748e887e 100755 --- a/structs.go +++ b/structs.go @@ -513,6 +513,7 @@ type ExecutionInfo struct { LastMonthlyResetMonth int `json:"last_monthly_reset_month" datastore:"last_monthly_reset_month"` LastUsageAlertThreshold int64 `json:"last_usage_alert_threshold" datastore:"last_usage_alert_threshold"` UsageAlerts []AlertThreshold `json:"usage_alerts" datastore:"usage_alerts"` + MonthlyAIUsageAlertSent bool `json:"monthly_ai_usage_alert_sent" datastore:"monthly_ai_usage_alert_sent"` } type AdditionalUseConfig struct {