Optimised the finaliser for AI Agents

Author: frikky

ai.go

@@ -6645,14 +6645,15 @@ SINGUL ACTIONS:
 		extraString = ""
 	}
 
+
 	// The starting decision number
 	lastFinishedIndex := -1
 
 	oldActionResult := ActionResult{}
 	_ = oldActionResult 
 	oldAgentOutput := AgentOutput{}
 	if createNextActions == true { 
-		extraString = "This is a continuation of a previous execution. ONLY output decisions that fit AFTER the last FINISHED decision. DO NOT repeat previous decisions, and make sure your indexing is on point. Output as an array of decisions.\n\nIF you don't want to add any new decision, add AT LEAST one decision saying why it is finished, summarising EXACTLY what the user wants in a user-friendly Markdown format, OR the format the user asked for. Make the action and category 'finish', and put the reason in the 'reason' field in a user friendly format."
+		extraString = "This is a continuation of a previous execution. ONLY output decisions that fit AFTER the last FINISHED decision. DO NOT repeat previous decisions, and make sure your indexing is on point. Output as an array of decisions.\n\nIF you don't want to add any new decision, add AT LEAST one decision saying why it is finished, summarising EXACTLY what the user wants in a user-friendly Markdown format, OR the format the user asked for. Make the action and category 'finish', and put the reason in the 'reason' field. Do NOT summarize or explain. JUST give exactly the final answer and nothing more."
 
 		userMessageChanged := false
 
@@ -6716,6 +6717,10 @@ SINGUL ACTIONS:
 				break
 			}
 
+			if len(userMessage) == 0 && len(oldAgentOutput.OriginalInput) > 0 {
+				userMessage = fmt.Sprintf("Original input: %s", oldAgentOutput.OriginalInput)
+			}
+
 			if hasFailure { 
 				userMessage += "\n\nSome of the previous decisions failed. Finalise the agent.\n\n"
 			}
@@ -6724,6 +6729,7 @@ SINGUL ACTIONS:
 			if len(previousAnswers) > 0 {
 				userMessage += fmt.Sprintf("\n\nAnswers to questions:\n%s", previousAnswers)
 			}
+
 			userMessage += "\n\nBased on the previous decisions, find out if any new decisions need to be added."
 			userMessageChanged = true
 		}
@@ -7283,7 +7289,6 @@ FINALISING:
 				}
 
 				SetWorkflowExecution(ctx, execution, true)
-				//os.Exit(3)
 			}
 		}
 
@@ -7301,6 +7306,8 @@ FINALISING:
 			// Ensures we track them along the way
 			if len(parsedAgentInput) > 0 { 
 				agentOutput.Input = parsedAgentInput
+
+				agentOutput.OriginalInput = userMessage
 			}
 		}
 
@@ -7347,16 +7354,30 @@ FINALISING:
 
 
 				agentOutput.Output = decision.Reason
+				for _, decisionField := range decision.Fields {
+					if (decisionField.Key == "output" || decisionField.Key == "body") && len(decisionField.Value) > 0 {
+						agentOutput.Output = decisionField.Value
+					} 
+				}
+
 				agentOutput.Status = "FINISHED"
 				agentOutput.CompletedAt = time.Now().Unix()
 
 				//workflowExecution.Results[resultIndex].Status = "SUCCESS"
 				//go sendAgentActionSelfRequest("SUCCESS", workflowExecution, workflowExecution.Results[resultIndex])
 
+			//} else if decision.Action == "answer" {
+			//	agentOutput.Decisions[decisionIndex].RunDetails.StartedAt = time.Now().Unix()
+			//	agentOutput.Decisions[decisionIndex].RunDetails.CompletedAt = time.Now().Unix()
+			//	agentOutput.Decisions[decisionIndex].RunDetails.Status = "FINISHED"
+
+				//go RunAgentDecisionAction(execution, agentOutput, agentOutput.Decisions[decisionIndex])
+
 			} else if decision.Action == "ask" || decision.Action == "question" { 
 				agentOutput.Decisions[decisionIndex].RunDetails.StartedAt = time.Now().Unix()
 				agentOutput.Decisions[decisionIndex].RunDetails.Status = "RUNNING"
 
+
 			} else if decision.Category != "standalone" { 
 				// Do we run the singul action directly?
 				agentOutput.Decisions[decisionIndex].RunDetails.StartedAt = time.Now().Unix()
@@ -7366,12 +7387,15 @@ FINALISING:
 
 
 			} else {
-				if decision.Category == "standalone" { 
 
+				if decision.Category == "standalone" || decision.Action == "answer" { 
 					// FIXME: Maybe need to send this to myself
-					decision.RunDetails.StartedAt = time.Now().Unix()
-					decision.RunDetails.CompletedAt = time.Now().Unix()
-					decision.RunDetails.Status = "FINISHED"
+
+					agentOutput.Decisions[decisionIndex].RunDetails.StartedAt = time.Now().Unix()
+					agentOutput.Decisions[decisionIndex].RunDetails.CompletedAt = time.Now().Unix()
+					agentOutput.Decisions[decisionIndex].RunDetails.Status = "FINISHED"
+
+					decision = agentOutput.Decisions[decisionIndex]
 
 					marshalledDecision, err := json.Marshal(decision)
 					if err != nil {
@@ -7391,7 +7415,7 @@ FINALISING:
 							Result: string(marshalledDecision),
 						}
 
-						// This is required as the result for the agent isn't set yet on the first run
+						// This is required as the result for the agent isn't set yet on the first run. Minor delay to wait up a bit
 						if decisionIndex == 0 {
 							go func() {
 								time.Sleep(2 * time.Second)
@@ -7449,7 +7473,7 @@ FINALISING:
 			// These aren't properly being updated in the db, so 
 			// we need additional logic here to ensure it is being 
 			// set/started
-			if nextActionType == "ask" || nextActionType == "finish" { 
+			if nextActionType == "ask" || nextActionType == "finish" || nextActionType == "answer" {  
 				// Ensure we update all of it
 				for resultIndex, result := range execution.Results {
 					if result.Action.ID != startNode.ID {

db-connector.go

@@ -1233,12 +1233,12 @@ func Fixexecution(ctx context.Context, workflowExecution WorkflowExecution) (Wor
 				continue
 			}
 
-			if innerresult.Status != "WAITING" {
+			if innerresult.Status != "WAITING" && innerresult.Status != "SUCCESS" {
 				found = true
 				result = innerresult
 				break
 
-			} else if innerresult.Status == "WAITING" && (action.AppName == "AI Agent" || action.AppName == "Shuffle Agent") {
+			} else if innerresult.Status == "WAITING" || innerresult.Status == "SUCCESS" && (action.AppName == "AI Agent" || action.AppName == "Shuffle Agent") {
 				// Auto fixing decision data based on cache for better decisionmaking
 				//log.Printf("[DEBUG] Found action result %s with WAITING status", action.AppName)
 
@@ -1253,23 +1253,44 @@ func Fixexecution(ctx context.Context, workflowExecution WorkflowExecution) (Wor
 
 				finishedDecisions := []string{}
 				failedFound := false
+				finishDecisionFound := false
 				// FIXME: Optimize it to not run too far past "FINISHED" on cache searches
 				for decisionIndex, decision := range mappedOutput.Decisions {
+					if decision.Action == "finish" {
+						finishDecisionFound = true
+					}
+
 					if decision.RunDetails.Status == "NOT IMPLEMENTED" {
 						continue
 					}
 
+					decisionId := fmt.Sprintf("agent-%s-%s", workflowExecution.ExecutionId, decision.RunDetails.Id)
 					if decision.RunDetails.Status == "FINISHED" {
 						finishedDecisions = append(finishedDecisions, decision.RunDetails.Id)
 						continue
 					} else if decision.RunDetails.Status == "FAILURE" {
 						//finishedDecisions = append(finishedDecisions, decision.RunDetails.Id)
 						failedFound = true
 						continue
+					} else {
+						if decision.RunDetails.CompletedAt > 0 {
+							if debug { 
+								log.Printf("[DEBUG] Rewriting decision %s to FINISHED based on completed at timestamp.", decision.RunDetails.Id)
+							}
+
+							mappedOutput.Decisions[decisionIndex].RunDetails.Status = "FINISHED"
+							finishedDecisions = append(finishedDecisions, decision.RunDetails.Id)
+							decisionsUpdated = true
+
+							marshalledDecision, err := json.Marshal(mappedOutput.Decisions[decisionIndex])
+							if err == nil {
+								err = SetCache(ctx, decisionId, marshalledDecision, 60)
+							}
+							continue
+						}
 					}
 
 					//log.Printf("[DEBUG] Check cache for %s with status %s", decision.RunDetails.Id, decision.RunDetails.Status)
-					decisionId := fmt.Sprintf("agent-%s-%s", workflowExecution.ExecutionId, decision.RunDetails.Id)
 					cache, err := GetCache(ctx, decisionId)
 					if err == nil {
 						foundDecision := AgentDecision{}
@@ -1303,12 +1324,23 @@ func Fixexecution(ctx context.Context, workflowExecution WorkflowExecution) (Wor
 				} 
 
 				if len(finishedDecisions) == len(mappedOutput.Decisions) && mappedOutput.Status != "FINISHED" && mappedOutput.Status != "FAILURE" && mappedOutput.Status != "ABORTED" {
-					decisionsUpdated = true
-					mappedOutput.Status = "FINISHED"
-					mappedOutput.CompletedAt = time.Now().Unix()
-					workflowExecution.Results[resultIndex].Status = "SUCCESS"
+					if finishDecisionFound {
+						decisionsUpdated = true
+						mappedOutput.Status = "FINISHED"
+						mappedOutput.CompletedAt = time.Now().Unix()
+
+						workflowExecution.Results[resultIndex].Status = "SUCCESS"
+						go sendAgentActionSelfRequest("SUCCESS", workflowExecution, workflowExecution.Results[resultIndex])
+					} else {
+						mappedOutput.Status = "WAITING"
+						mappedOutput.CompletedAt = 0
+						go sendAgentActionSelfRequest("SUCCESS", workflowExecution, workflowExecution.Results[resultIndex])
+
+						if debug { 
+							log.Printf("[DEBUG] All decisions finished but no finish action found.")
+						}
+					}
 
-					go sendAgentActionSelfRequest("SUCCESS", workflowExecution, workflowExecution.Results[resultIndex])
 				}
 
 				if decisionsUpdated {

structs.go

@@ -1705,9 +1705,9 @@ type SelectedDetectionRules struct {
 }
 
 type DetectionFileInfo struct {
-	FileName string `json:"file_name" yaml:"file_name"`
-	FileId   string `json:"file_id"`
-	Tags 	[]string `json:"tags" yaml:"tags"`
+	FileName string   `json:"file_name" yaml:"file_name"`
+	FileId   string   `json:"file_id"`
+	Tags     []string `json:"tags" yaml:"tags"`
 
 	RuleTitle   string `json:"title" yaml:"title"`
 	Description string `json:"description" yaml:"description"`
@@ -4578,13 +4578,14 @@ type AgentOutput struct {
 	// For easy testing
 	DecisionString string `json:"decision_string,omitempty" datastore:"decision_string"`
 	// For tracking of details parent<->child
-	StartedAt   int64  `json:"started_at,omitempty" datastore:"started_at"`
-	CompletedAt int64  `json:"completed_at,omitempty" datastore:"completed_at"`
-	ExecutionId string `json:"execution_id,omitempty" datastore:"execution_id"`
-	NodeId      string `json:"node_id,omitempty" datastore:"node_id"`
-	Memory      string `json:"memory,omitempty" datastore:"memory"`
-	Input       string `json:"input" datastore:"input"`
-	Output      string `json:"output,omitempty" datastore:"output"`
+	StartedAt     int64  `json:"started_at,omitempty" datastore:"started_at"`
+	CompletedAt   int64  `json:"completed_at,omitempty" datastore:"completed_at"`
+	ExecutionId   string `json:"execution_id,omitempty" datastore:"execution_id"`
+	NodeId        string `json:"node_id,omitempty" datastore:"node_id"`
+	Memory        string `json:"memory,omitempty" datastore:"memory"`
+	Input         string `json:"input" datastore:"input"`
+	Output        string `json:"output,omitempty" datastore:"output"`
+	OriginalInput string `json:"original_input,omitempty" datastore:"original_input"`
 }
 
 type HTTPWrapper struct {
@@ -4749,8 +4750,8 @@ type ProcessInfo struct {
 }
 
 type UserInfo struct {
-	UserID   string `json:"user_id"`
-	Username string `json:"username"`
+	UserID   string   `json:"user_id"`
+	Username string   `json:"username"`
 	Groups   []string `json:"groups,omitempty"`
 }
 
@@ -4763,9 +4764,9 @@ type TelemetryConfig struct {
 }
 
 type TelemetryFilter struct {
-	Type     string   `json:"type"`
-	Include  []string `json:"include,omitempty"`
-	Exclude  []string `json:"exclude,omitempty"`
+	Type    string   `json:"type"`
+	Include []string `json:"include,omitempty"`
+	Exclude []string `json:"exclude,omitempty"`
 }
 
 type AuditLogCollector struct {