63 lines
2.4 KiB
Go
63 lines
2.4 KiB
Go
package config
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"time"
|
|
|
|
"github.com/spf13/viper"
|
|
)
|
|
|
|
// Config holds the benchmark configuration.
|
|
type Config struct {
|
|
APIEndpoint string `yaml:"apiEndpoint" mapstructure:"apiEndpoint"`
|
|
APIKey string `yaml:"apiKey" mapstructure:"apiKey"` // Consider loading from env or secure store
|
|
Model string `yaml:"model" mapstructure:"model"`
|
|
Requests int `yaml:"requests"`
|
|
Concurrency int `yaml:"concurrency"`
|
|
RateLimit int `yaml:"rateLimit"` // Requests per second
|
|
Duration time.Duration `yaml:"duration"`
|
|
PromptTokens int `yaml:"promptTokens"`
|
|
MaxTokens int `yaml:"maxTokens"`
|
|
Streaming bool `yaml:"streaming" mapstructure:"streaming"`
|
|
Timeout time.Duration `yaml:"timeout"` // Request timeout
|
|
RequestTimeout time.Duration `yaml:"requestTimeout" mapstructure:"request_timeout"` // Timeout for individual HTTP requests (defaults to global Timeout if 0)
|
|
OutputReport string `yaml:"outputReport"` // Path to save the HTML report
|
|
Headers map[string]string `yaml:"headers"` // Custom HTTP headers
|
|
Payload map[string]any `yaml:"payload"` // Custom payload fields
|
|
Client string `yaml:"client"` // "fasthttp" or "nethttp"
|
|
ReportFile string `yaml:"reportFile" mapstructure:"report_file"` // Output file for the HTML report
|
|
// Add other configuration parameters as needed
|
|
}
|
|
|
|
// LoadConfig loads configuration from the specified file path.
|
|
func LoadConfig(filePath string) (*Config, error) {
|
|
// Check if the file exists
|
|
if _, err := os.Stat(filePath); os.IsNotExist(err) {
|
|
return nil, fmt.Errorf("config file not found at path: %s", filePath)
|
|
}
|
|
|
|
v := viper.New()
|
|
v.SetConfigFile(filePath)
|
|
v.SetConfigType("yaml")
|
|
|
|
// Read the configuration file
|
|
if err := v.ReadInConfig(); err != nil {
|
|
return nil, fmt.Errorf("failed to read config file '%s': %w", filePath, err)
|
|
}
|
|
|
|
var config Config
|
|
// Unmarshal the configuration into the Config struct
|
|
if err := v.Unmarshal(&config); err != nil {
|
|
return nil, fmt.Errorf("failed to unmarshal config from '%s': %w", filePath, err)
|
|
}
|
|
|
|
// Set default RequestTimeout if not provided
|
|
if config.RequestTimeout <= 0 {
|
|
config.RequestTimeout = config.Timeout
|
|
}
|
|
|
|
// TODO: Add validation for required fields if necessary
|
|
return &config, nil
|
|
}
|