package test import ( "fmt" "net/url" "strconv" "strings" jsoniter "github.com/json-iterator/go" goutext "github.com/yaoapp/gou/text" "github.com/yaoapp/yao/agent/assistant" "github.com/yaoapp/yao/agent/context" ) // InputSourceType represents the type of input source type InputSourceType string const ( // InputSourceFile indicates input from a JSONL file InputSourceFile InputSourceType = "file" // InputSourceMessage indicates input from a direct message string InputSourceMessage InputSourceType = "message" // InputSourceScript indicates script test mode InputSourceScript InputSourceType = "script" // InputSourceAgent indicates input generated by an agent InputSourceAgent InputSourceType = "agent" ) // InputSource represents a parsed input source type InputSource struct { Type InputSourceType // file, message, script, agent Value string // path, message, script ref, or agent ID Params map[string]interface{} // query parameters (for agent source) } // ParseInputSource parses the -i flag value into an InputSource // Supported formats: // - "agents:workers.test.generator" - Agent-generated test cases // - "agents:workers.test.generator?count=10&focus=edge-cases" - With parameters // - "scripts.tests.gen" - Script-generated test cases // - "./tests/inputs.jsonl" - JSONL file // - "Hello, how are you?" - Direct message func ParseInputSource(input string) *InputSource { // Check for agents: prefix if strings.HasPrefix(input, "agents:") { return parseAgentSource(strings.TrimPrefix(input, "agents:")) } // Check for scripts: prefix (for generator scripts) if strings.HasPrefix(input, "scripts:") { return &InputSource{ Type: InputSourceScript, Value: strings.TrimPrefix(input, "scripts:"), } } // Check for script test mode (scripts.xxx format without prefix) if strings.HasPrefix(input, "scripts.") { return &InputSource{ Type: InputSourceScript, Value: input, } } // Check for file extension if strings.HasSuffix(input, ".jsonl") || strings.HasSuffix(input, ".json") { return &InputSource{ Type: InputSourceFile, Value: input, } } // Check if it looks like a file path if strings.Contains(input, "/") || strings.Contains(input, "\\") { return &InputSource{ Type: InputSourceFile, Value: input, } } // Default to message return &InputSource{ Type: InputSourceMessage, Value: input, } } // parseAgentSource parses an agent source string with optional query parameters // Format: "agent.id" or "agent.id?count=10&focus=edge-cases" func parseAgentSource(input string) *InputSource { source := &InputSource{ Type: InputSourceAgent, Params: make(map[string]interface{}), } // Check for query parameters if idx := strings.Index(input, "?"); idx >= 0 { source.Value = input[:idx] queryStr := input[idx+1:] // Parse query parameters values, err := url.ParseQuery(queryStr) if err == nil { for key, vals := range values { if len(vals) > 0 { // Try to parse as number if num, err := strconv.Atoi(vals[0]); err == nil { source.Params[key] = num } else if num, err := strconv.ParseFloat(vals[0], 64); err == nil { source.Params[key] = num } else if vals[0] == "true" { source.Params[key] = true } else if vals[0] == "false" { source.Params[key] = false } else { source.Params[key] = vals[0] } } } } } else { source.Value = input } return source } // GeneratorInput represents the input sent to a generator agent type GeneratorInput struct { TargetAgent *TargetAgentInfo `json:"target_agent"` Count int `json:"count,omitempty"` Focus string `json:"focus,omitempty"` Extra map[string]interface{} `json:"extra,omitempty"` } // TargetAgentInfo contains information about the agent being tested type TargetAgentInfo struct { ID string `json:"id"` Description string `json:"description,omitempty"` Tools []map[string]interface{} `json:"tools,omitempty"` } // GenerateTestCases generates test cases using a generator agent func GenerateTestCases(agentID string, targetInfo *TargetAgentInfo, params map[string]interface{}) ([]*Case, error) { // Get generator assistant ast, err := assistant.Get(agentID) if err != nil { return nil, fmt.Errorf("failed to get generator agent %s: %w", agentID, err) } // Build generation request genInput := &GeneratorInput{ TargetAgent: targetInfo, Count: 5, // Default count } // Apply parameters if params != nil { if count, ok := params["count"].(int); ok { genInput.Count = count } if focus, ok := params["focus"].(string); ok { genInput.Focus = focus } // Store extra parameters genInput.Extra = make(map[string]interface{}) for k, v := range params { if k != "count" && k != "focus" { genInput.Extra[k] = v } } } // Create context env := NewEnvironment("", "") ctx := NewTestContext("generator", agentID, env) defer ctx.Release() // Build options - skip history and trace for efficiency opts := &context.Options{ Skip: &context.Skip{ History: true, Trace: true, Output: true, }, Metadata: map[string]interface{}{ "test_mode": "generator", }, } // Build message inputJSON, err := jsoniter.Marshal(genInput) if err != nil { return nil, fmt.Errorf("failed to marshal generator input: %w", err) } messages := []context.Message{{ Role: context.RoleUser, Content: string(inputJSON), }} // Call generator agent response, err := ast.Stream(ctx, messages, opts) if err != nil { return nil, fmt.Errorf("generator agent error: %w", err) } // Extract and parse response return parseGeneratedCases(response) } // parseGeneratedCases parses the generator agent's response into test cases func parseGeneratedCases(response *context.Response) ([]*Case, error) { if response == nil || response.Completion == nil { return nil, fmt.Errorf("empty response from generator agent") } // Extract content content := response.Completion.Content if content == nil { return nil, fmt.Errorf("no content in generator response") } // Convert content to string var text string switch v := content.(type) { case string: text = v default: data, err := jsoniter.Marshal(content) if err != nil { return nil, fmt.Errorf("failed to marshal content: %w", err) } text = string(data) } // Use goutext.ExtractJSON for fault-tolerant parsing parsed := goutext.ExtractJSON(text) if parsed == nil { return nil, fmt.Errorf("failed to parse generator response as JSON: %s", truncateOutput(text, 200)) } // Convert to []*Case return convertToCases(parsed) } // convertToCases converts parsed JSON to test cases func convertToCases(parsed interface{}) ([]*Case, error) { // Handle array of cases arr, ok := parsed.([]interface{}) if !ok { // Maybe it's a single case wrapped in an object if obj, ok := parsed.(map[string]interface{}); ok { if cases, ok := obj["cases"].([]interface{}); ok { arr = cases } else if testCases, ok := obj["test_cases"].([]interface{}); ok { arr = testCases } else { // Single case arr = []interface{}{obj} } } else { return nil, fmt.Errorf("expected array of test cases, got %T", parsed) } } cases := make([]*Case, 0, len(arr)) for i, item := range arr { caseMap, ok := item.(map[string]interface{}) if !ok { return nil, fmt.Errorf("test case %d is not an object", i) } tc, err := mapToCase(caseMap) if err != nil { return nil, fmt.Errorf("failed to parse test case %d: %w", i, err) } cases = append(cases, tc) } return cases, nil } // mapToCase converts a map to a Case struct func mapToCase(m map[string]interface{}) (*Case, error) { tc := &Case{} // Required: id if id, ok := m["id"].(string); ok { tc.ID = id } else { return nil, fmt.Errorf("missing required field 'id'") } // Required: input if input, ok := m["input"]; ok { tc.Input = input } else { return nil, fmt.Errorf("missing required field 'input'") } // Optional: assertions/assert if assertions, ok := m["assertions"]; ok { tc.Assert = assertions } else if assert, ok := m["assert"]; ok { tc.Assert = assert } // Optional: options - convert map to CaseOptions if options, ok := m["options"].(map[string]interface{}); ok { tc.Options = mapToCaseOptions(options) } // Optional: before/after if before, ok := m["before"].(string); ok { tc.Before = before } if after, ok := m["after"].(string); ok { tc.After = after } // Optional: timeout if timeout, ok := m["timeout"].(string); ok { tc.Timeout = timeout } return tc, nil } // ToInputMode converts InputSourceType to InputMode for backward compatibility func (s *InputSource) ToInputMode() InputMode { switch s.Type { case InputSourceFile: return InputModeFile case InputSourceMessage: return InputModeMessage case InputSourceScript: return InputModeScript case InputSourceAgent: // Agent source generates cases, then runs in file mode return InputModeFile default: return InputModeMessage } } // mapToCaseOptions converts a map to CaseOptions func mapToCaseOptions(m map[string]interface{}) *CaseOptions { opts := &CaseOptions{} if connector, ok := m["connector"].(string); ok { opts.Connector = connector } if mode, ok := m["mode"].(string); ok { opts.Mode = mode } if disableGlobalPrompts, ok := m["disable_global_prompts"].(bool); ok { opts.DisableGlobalPrompts = disableGlobalPrompts } if search, ok := m["search"].(bool); ok { opts.Search = &search } if metadata, ok := m["metadata"].(map[string]interface{}); ok { opts.Metadata = metadata } if skip, ok := m["skip"].(map[string]interface{}); ok { opts.Skip = &CaseSkipOptions{} if history, ok := skip["history"].(bool); ok { opts.Skip.History = history } if trace, ok := skip["trace"].(bool); ok { opts.Skip.Trace = trace } if output, ok := skip["output"].(bool); ok { opts.Skip.Output = output } if keyword, ok := skip["keyword"].(bool); ok { opts.Skip.Keyword = keyword } if searchSkip, ok := skip["search"].(bool); ok { opts.Skip.Search = searchSkip } } return opts }