picoclaw/pkg/evolution/success_judge.go

package evolution

import (
	"context"
	"encoding/json"
	"strings"

	"github.com/sipeed/picoclaw/pkg/providers"
)

type TaskSuccessDecision struct {
	Success bool
	Reason  string
}

type SuccessJudge interface {
	JudgeTaskRecord(ctx context.Context, record LearningRecord) (TaskSuccessDecision, error)
}

type HeuristicSuccessJudge struct{}

func (j *HeuristicSuccessJudge) JudgeTaskRecord(
	_ context.Context,
	record LearningRecord,
) (TaskSuccessDecision, error) {
	if record.Success == nil || !*record.Success {
		return TaskSuccessDecision{Success: false, Reason: "task not completed"}, nil
	}
	if strings.TrimSpace(record.UserGoal) == "" {
		return TaskSuccessDecision{Success: false, Reason: "missing user goal"}, nil
	}
	if strings.EqualFold(strings.TrimSpace(record.SessionKey), "heartbeat") {
		return TaskSuccessDecision{Success: false, Reason: "heartbeat session"}, nil
	}
	if strings.EqualFold(strings.TrimSpace(record.FinalOutput), "HEARTBEAT_OK") {
		return TaskSuccessDecision{Success: false, Reason: "heartbeat output"}, nil
	}
	if len(record.ToolExecutions) > 0 {
		allFailed := true
		for _, exec := range record.ToolExecutions {
			if exec.Success {
				allFailed = false
				break
			}
		}
		if allFailed {
			return TaskSuccessDecision{Success: false, Reason: "all tools failed"}, nil
		}
	}
	if strings.TrimSpace(record.FinalOutput) == "" && len(record.ToolExecutions) == 0 {
		return TaskSuccessDecision{Success: false, Reason: "missing final output"}, nil
	}
	return TaskSuccessDecision{Success: true, Reason: "heuristic success"}, nil
}

type LLMTaskSuccessJudge struct {
	provider providers.LLMProvider
	model    string
	fallback SuccessJudge
}

type llmTaskSuccessResponse struct {
	Success bool   `json:"success"`
	Reason  string `json:"reason"`
}

func NewLLMTaskSuccessJudge(provider providers.LLMProvider, model string, fallback SuccessJudge) *LLMTaskSuccessJudge {
	if fallback == nil {
		fallback = &HeuristicSuccessJudge{}
	}
	return &LLMTaskSuccessJudge{
		provider: provider,
		model:    strings.TrimSpace(model),
		fallback: fallback,
	}
}

func (j *LLMTaskSuccessJudge) JudgeTaskRecord(
	ctx context.Context,
	record LearningRecord,
) (TaskSuccessDecision, error) {
	if j == nil || j.provider == nil {
		return j.fallbackDecision(ctx, record)
	}

	model := strings.TrimSpace(j.model)
	if model == "" {
		model = strings.TrimSpace(j.provider.GetDefaultModel())
	}
	if model == "" {
		return j.fallbackDecision(ctx, record)
	}

	resp, err := j.provider.Chat(ctx, []providers.Message{
		{
			Role:    "system",
			Content: "Return exactly one JSON object with fields success:boolean and reason:string. No markdown fences.",
		},
		{
			Role:    "user",
			Content: buildTaskSuccessJudgePrompt(record),
		},
	}, nil, model, map[string]any{"temperature": 0})
	if err != nil || resp == nil {
		return j.fallbackDecision(ctx, record)
	}

	content := strings.TrimSpace(resp.Content)
	content = strings.TrimPrefix(content, "```json")
	content = strings.TrimPrefix(content, "```")
	content = strings.TrimSuffix(content, "```")
	content = strings.TrimSpace(content)
	if content == "" {
		return j.fallbackDecision(ctx, record)
	}

	var payload llmTaskSuccessResponse
	if err := json.Unmarshal([]byte(content), &payload); err != nil {
		return j.fallbackDecision(ctx, record)
	}
	return TaskSuccessDecision{
		Success: payload.Success,
		Reason:  strings.TrimSpace(payload.Reason),
	}, nil
}

func (j *LLMTaskSuccessJudge) fallbackDecision(
	ctx context.Context,
	record LearningRecord,
) (TaskSuccessDecision, error) {
	if j == nil || j.fallback == nil {
		return TaskSuccessDecision{Success: false, Reason: "no success judge available"}, nil
	}
	return j.fallback.JudgeTaskRecord(ctx, record)
}

func buildTaskSuccessJudgePrompt(record LearningRecord) string {
	lines := []string{
		"Decide whether this agent task truly achieved the user's goal.",
		"Reject tasks that are only partial reasoning, only describe future steps, or obviously did not complete the requested outcome.",
		"",
		"User goal: " + fallbackString(record.UserGoal, "none"),
		"Summary: " + fallbackString(record.Summary, "none"),
		"Final output: " + fallbackString(record.FinalOutput, "none"),
		"Used skills: " + joinOrFallback(record.UsedSkillNames, "none"),
		"Tool kinds: " + joinOrFallback(record.ToolKinds, "none"),
		"Tool executions:",
	}
	if len(record.ToolExecutions) == 0 {
		lines = append(lines, "- none")
	} else {
		for _, exec := range record.ToolExecutions {
			status := "failed"
			if exec.Success {
				status = "succeeded"
			}
			line := "- " + exec.Name + ": " + status
			if errSummary := strings.TrimSpace(exec.ErrorSummary); errSummary != "" {
				line += " (" + errSummary + ")"
			}
			if len(exec.SkillNames) > 0 {
				line += " [skills: " + strings.Join(exec.SkillNames, ", ") + "]"
			}
			lines = append(lines, line)
		}
	}
	return strings.Join(lines, "\n")
}