Added averaging of three runs, Refactor testProviderMetrics for improved error handling and logging; enhance metric calculations and reporting for better clarity and accuracy

Author: lemon07r

main.go

@@ -40,34 +40,19 @@ type TestResult struct {
 	Error            string        `json:"error,omitempty"`
 }
 
-// testProviderMetrics runs a full benchmark test against a single provider.
-// It is designed to be run as a goroutine.
-func testProviderMetrics(config ProviderConfig, tke *tiktoken.Tiktoken, wg *sync.WaitGroup, logDir, resultsDir string, results *[]TestResult, resultsMutex *sync.Mutex) {
-	// Defer wg.Done() if this is part of a concurrent group
-	if wg != nil {
-		defer wg.Done()
-	}
-
-	// Create log file for this provider
-	timestamp := time.Now().Format("20060102-150405")
-	logFile, err := os.Create(filepath.Join(logDir, fmt.Sprintf("%s-%s.log", config.Name, timestamp)))
-	if err != nil {
-		log.Printf("Error creating log file for %s: %v", config.Name, err)
-		return
-	}
-	defer logFile.Close()
-
-	// Create a logger for this provider that writes to both stdout and file
-	providerLogger := log.New(io.MultiWriter(os.Stdout, logFile), "", log.LstdFlags)
-
-	providerLogger.Printf("--- Testing: %s (%s) ---", config.Name, config.Model)
+// formatDuration formats a duration as decimal seconds
+func formatDuration(d time.Duration) string {
+	return fmt.Sprintf("%.3fs", d.Seconds())
+}
 
-	// 5. Configure the OpenAI Client
+// singleTestRun performs one test run and returns metrics or error
+func singleTestRun(config ProviderConfig, tke *tiktoken.Tiktoken, providerLogger *log.Logger, ctx context.Context) (e2e, ttft time.Duration, throughput float64, tokens int, err error) {
+	// Configure the OpenAI Client
 	clientConfig := openai.DefaultConfig(config.APIKey)
 	clientConfig.BaseURL = config.BaseURL
 	client := openai.NewClientWithConfig(clientConfig)
 
-	// 6. Define the request
+	// Define the request
 	prompt := "You are a helpful assistant. Please write a short, 150-word story about a curious robot exploring an ancient, overgrown library on a forgotten planet."
 	messages := []openai.ChatCompletionMessage{
 		{
@@ -83,65 +68,37 @@ func testProviderMetrics(config ProviderConfig, tke *tiktoken.Tiktoken, wg *sync
 		Stream:    true,
 	}
 
-	// 7. Execute the stream and measure metrics
-	startTime := time.Now() // ---- START TIMER
+	// Execute the stream and measure metrics
+	startTime := time.Now()
 	var firstTokenTime time.Time
 	var fullResponseContent strings.Builder
 
-	// Add timeout context to prevent indefinite hangs (2 minutes)
-	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
-	defer cancel()
-
-	stream, err := client.CreateChatCompletionStream(ctx, req)
-	if err != nil {
-		providerLogger.Printf("Error creating stream for %s: %v", config.Name, err)
-		// Save error result
-		result := TestResult{
-			Provider:  config.Name,
-			Model:     config.Model,
-			Timestamp: time.Now(),
-			Success:   false,
-			Error:     err.Error(),
-		}
-		saveResult(resultsDir, result)
-		appendResult(results, resultsMutex, result)
-		return
+	stream, streamErr := client.CreateChatCompletionStream(ctx, req)
+	if streamErr != nil {
+		return 0, 0, 0, 0, fmt.Errorf("error creating stream: %w", streamErr)
 	}
-	defer stream.Close() // IMPORTANT: Always close the stream
+	defer stream.Close()
 
 	providerLogger.Printf("[%s] ... Request sent. Waiting for stream ...", config.Name)
 
 	for {
-		response, err := stream.Recv()
+		response, recvErr := stream.Recv()
 
 		// Check for end of stream
-		if errors.Is(err, io.EOF) {
+		if errors.Is(recvErr, io.EOF) {
 			providerLogger.Printf("[%s] ... Stream complete.", config.Name)
 			break
 		}
 
-		if err != nil {
-			errMsg := err.Error()
+		if recvErr != nil {
 			if ctx.Err() == context.DeadlineExceeded {
-				errMsg = "Timeout: stream took longer than 5 minutes"
+				return 0, 0, 0, 0, fmt.Errorf("timeout exceeded")
 			}
-			providerLogger.Printf("Stream error for %s: %v", config.Name, errMsg)
-			// Save error result
-			result := TestResult{
-				Provider:  config.Name,
-				Model:     config.Model,
-				Timestamp: time.Now(),
-				Success:   false,
-				Error:     errMsg,
-			}
-			saveResult(resultsDir, result)
-			appendResult(results, resultsMutex, result)
-			return
+			return 0, 0, 0, 0, fmt.Errorf("stream error: %w", recvErr)
 		}
 
-		// Check if Choices array is empty (some APIs send empty chunks)
+		// Check if Choices array is empty
 		if len(response.Choices) == 0 {
-			providerLogger.Printf("[%s] ... Received empty chunk (no Choices)", config.Name)
 			continue
 		}
 
@@ -150,7 +107,7 @@ func testProviderMetrics(config ProviderConfig, tke *tiktoken.Tiktoken, wg *sync
 
 		// Check if this is the first chunk with actual text
 		if content != "" && firstTokenTime.IsZero() {
-			firstTokenTime = time.Now() // ---- TTFT METRIC
+			firstTokenTime = time.Now()
 			providerLogger.Printf("[%s] ... First token received!", config.Name)
 		}
 
@@ -160,23 +117,10 @@ func testProviderMetrics(config ProviderConfig, tke *tiktoken.Tiktoken, wg *sync
 		}
 	}
 
-	endTime := time.Now() // ---- E2E METRIC
-
-	// --- 8. Calculate and Print Results ---
+	endTime := time.Now()
 
 	if firstTokenTime.IsZero() {
-		providerLogger.Printf("Error for %s: Did not receive any content from the API.", config.Name)
-		// Save error result
-		result := TestResult{
-			Provider:  config.Name,
-			Model:     config.Model,
-			Timestamp: time.Now(),
-			Success:   false,
-			Error:     "No content received from API",
-		}
-		saveResult(resultsDir, result)
-		appendResult(results, resultsMutex, result)
-		return
+		return 0, 0, 0, 0, fmt.Errorf("no content received from API")
 	}
 
 	// Get accurate token count
@@ -185,60 +129,125 @@ func testProviderMetrics(config ProviderConfig, tke *tiktoken.Tiktoken, wg *sync
 	completionTokens := len(tokenList)
 
 	if completionTokens == 0 {
-		providerLogger.Printf("Error for %s: Received response with 0 tokens.", config.Name)
-		// Save error result
-		result := TestResult{
-			Provider:  config.Name,
-			Model:     config.Model,
-			Timestamp: time.Now(),
-			Success:   false,
-			Error:     "Received 0 tokens",
-		}
-		saveResult(resultsDir, result)
-		appendResult(results, resultsMutex, result)
-		return
+		return 0, 0, 0, 0, fmt.Errorf("received 0 tokens")
 	}
 
-	// 1. End-to-End Latency
+	// Calculate metrics
 	e2eLatency := endTime.Sub(startTime)
+	ttftLatency := firstTokenTime.Sub(startTime)
+	generationTime := e2eLatency - ttftLatency
 
-	// 2. Time to First Token (TTFT)
-	ttft := firstTokenTime.Sub(startTime)
-
-	// 3. Throughput (Tokens per Second)
-	// This is (Total Tokens - 1) / (Time from first token to last token)
-	generationTime := e2eLatency - ttft
-	var throughput float64
-
+	var throughputVal float64
 	if generationTime.Seconds() <= 0 {
-		// Handle edge case where generation is too fast or only 1 token
-		throughput = 0.0
+		throughputVal = 0.0
 	} else {
-		throughput = (float64(completionTokens) - 1.0) / generationTime.Seconds()
+		throughputVal = (float64(completionTokens) - 1.0) / generationTime.Seconds()
+	}
+
+	return e2eLatency, ttftLatency, throughputVal, completionTokens, nil
+}
+
+// testProviderMetrics runs a full benchmark test against a single provider.
+// It runs 3 iterations and reports averaged results, with a 2-minute total timeout.
+func testProviderMetrics(config ProviderConfig, tke *tiktoken.Tiktoken, wg *sync.WaitGroup, logDir, resultsDir string, results *[]TestResult, resultsMutex *sync.Mutex) {
+	// Defer wg.Done() if this is part of a concurrent group
+	if wg != nil {
+		defer wg.Done()
 	}
 
-	// --- Print Results (use providerLogger for thread-safety) ---
+	// Create log file for this provider
+	timestamp := time.Now().Format("20060102-150405")
+	logFile, err := os.Create(filepath.Join(logDir, fmt.Sprintf("%s-%s.log", config.Name, timestamp)))
+	if err != nil {
+		log.Printf("Error creating log file for %s: %v", config.Name, err)
+		return
+	}
+	defer logFile.Close()
+
+	// Create a logger for this provider that writes to both stdout and file
+	providerLogger := log.New(io.MultiWriter(os.Stdout, logFile), "", log.LstdFlags)
+
+	providerLogger.Printf("--- Testing: %s (%s) - Running 3 iterations ---", config.Name, config.Model)
+
+	// Create 2-minute timeout context for all runs
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
+	defer cancel()
+
+	// Run up to 3 iterations and collect metrics
+	const maxIterations = 3
+	var e2eSum, ttftSum time.Duration
+	var throughputSum float64
+	var tokensSum int
+	successfulRuns := 0
+
+	for i := 1; i <= maxIterations; i++ {
+		// Check if timeout exceeded before starting next run
+		if ctx.Err() != nil {
+			providerLogger.Printf("[%s] Timeout reached after %d run(s)", config.Name, successfulRuns)
+			break
+		}
+
+		providerLogger.Printf("[%s] Run %d/%d", config.Name, i, maxIterations)
+
+		e2e, ttft, throughput, tokens, runErr := singleTestRun(config, tke, providerLogger, ctx)
+		if runErr != nil {
+			providerLogger.Printf("[%s] Run %d failed: %v", config.Name, i, runErr)
+			// If no successful runs yet, save error result
+			if successfulRuns == 0 && i == maxIterations {
+				result := TestResult{
+					Provider:  config.Name,
+					Model:     config.Model,
+					Timestamp: time.Now(),
+					Success:   false,
+					Error:     runErr.Error(),
+				}
+				saveResult(resultsDir, result)
+				appendResult(results, resultsMutex, result)
+			}
+			break
+		}
+
+		e2eSum += e2e
+		ttftSum += ttft
+		throughputSum += throughput
+		tokensSum += tokens
+		successfulRuns++
+
+		providerLogger.Printf("[%s] Run %d complete: E2E=%s TTFT=%s Throughput=%.2f tok/s",
+			config.Name, i, formatDuration(e2e), formatDuration(ttft), throughput)
+	}
+
+	if successfulRuns == 0 {
+		providerLogger.Printf("[%s] All runs failed", config.Name)
+		return
+	}
+
+	// Calculate averages
+	avgE2E := e2eSum / time.Duration(successfulRuns)
+	avgTTFT := ttftSum / time.Duration(successfulRuns)
+	avgThroughput := throughputSum / float64(successfulRuns)
+	avgTokens := tokensSum / successfulRuns
+
+	// Print averaged results
 	providerLogger.Println("==============================================")
-	providerLogger.Printf("   LLM Metrics for: %s", config.Name)
+	providerLogger.Printf("   LLM Metrics for: %s (averaged over %d run(s))", config.Name, successfulRuns)
 	providerLogger.Printf("   Model: %s", config.Model)
-	providerLogger.Printf("   Total Output Tokens: %d", completionTokens)
+	providerLogger.Printf("   Avg Output Tokens: %d", avgTokens)
 	providerLogger.Println("----------------------------------------------")
-	providerLogger.Printf("   End-to-End Latency: %v", e2eLatency)
-	providerLogger.Printf("   Latency (TTFT):     %v", ttft)
-	providerLogger.Printf("   Throughput (Tokens/sec): %.2f tokens/s", throughput)
+	providerLogger.Printf("   End-to-End Latency: %s", formatDuration(avgE2E))
+	providerLogger.Printf("   Latency (TTFT):     %s", formatDuration(avgTTFT))
+	providerLogger.Printf("   Throughput (Tokens/sec): %.2f tokens/s", avgThroughput)
 	providerLogger.Println("==============================================")
-	// Uncomment to see the full response
-	// providerLogger.Printf("[%s] Full Response:\n%s\n", config.Name, fullResponse)
 
 	// Save successful result
 	result := TestResult{
 		Provider:         config.Name,
 		Model:            config.Model,
 		Timestamp:        time.Now(),
-		E2ELatency:       e2eLatency,
-		TTFT:             ttft,
-		Throughput:       throughput,
-		CompletionTokens: completionTokens,
+		E2ELatency:       avgE2E,
+		TTFT:             avgTTFT,
+		Throughput:       avgThroughput,
+		CompletionTokens: avgTokens,
 		Success:          true,
 	}
 	saveResult(resultsDir, result)
@@ -306,11 +315,11 @@ func generateMarkdownReport(resultsDir string, results []TestResult, sessionTime
 
 		for _, r := range results {
 			if r.Success {
-				report.WriteString(fmt.Sprintf("| %s | %s | %v | %v | %.2f tok/s | %d |\n",
+				report.WriteString(fmt.Sprintf("| %s | %s | %s | %s | %.2f tok/s | %d |\n",
 					r.Provider,
 					r.Model,
-					r.E2ELatency,
-					r.TTFT,
+					formatDuration(r.E2ELatency),
+					formatDuration(r.TTFT),
 					r.Throughput,
 					r.CompletionTokens))
 			}
@@ -361,12 +370,12 @@ func generateMarkdownReport(resultsDir string, results []TestResult, sessionTime
 		report.WriteString("|------|----------|------------|------|-------------|\n")
 
 		for i, r := range successfulResults {
-			report.WriteString(fmt.Sprintf("| %d | %s | %.2f tok/s | %v | %v |\n",
+			report.WriteString(fmt.Sprintf("| %d | %s | %.2f tok/s | %s | %s |\n",
 				i+1,
 				r.Provider,
 				r.Throughput,
-				r.TTFT,
-				r.E2ELatency))
+				formatDuration(r.TTFT),
+				formatDuration(r.E2ELatency)))
 		}
 		report.WriteString("\n")
 
@@ -385,12 +394,12 @@ func generateMarkdownReport(resultsDir string, results []TestResult, sessionTime
 		report.WriteString("|------|----------|------|------------|-------------|\n")
 
 		for i, r := range successfulResults {
-			report.WriteString(fmt.Sprintf("| %d | %s | %v | %.2f tok/s | %v |\n",
+			report.WriteString(fmt.Sprintf("| %d | %s | %s | %.2f tok/s | %s |\n",
 				i+1,
 				r.Provider,
-				r.TTFT,
+				formatDuration(r.TTFT),
 				r.Throughput,
-				r.E2ELatency))
+				formatDuration(r.E2ELatency)))
 		}
 		report.WriteString("\n")
 
@@ -409,11 +418,11 @@ func generateMarkdownReport(resultsDir string, results []TestResult, sessionTime
 		report.WriteString("|------|----------|-------------|------|------------|\n")
 
 		for i, r := range successfulResults {
-			report.WriteString(fmt.Sprintf("| %d | %s | %v | %v | %.2f tok/s |\n",
+			report.WriteString(fmt.Sprintf("| %d | %s | %s | %s | %.2f tok/s |\n",
 				i+1,
 				r.Provider,
-				r.E2ELatency,
-				r.TTFT,
+				formatDuration(r.E2ELatency),
+				formatDuration(r.TTFT),
 				r.Throughput))
 		}
 		report.WriteString("\n")