57 lines
1.4 KiB
Go
57 lines
1.4 KiB
Go
package config
|
|
|
|
import (
|
|
"io/ioutil"
|
|
|
|
"gopkg.in/yaml.v3"
|
|
)
|
|
|
|
// Config 定义压测配置结构
|
|
// 可根据 production.md 扩展字段
|
|
|
|
type APIConfig struct {
|
|
Endpoint string `yaml:"endpoint"`
|
|
APIKey string `yaml:"api_key"`
|
|
Model string `yaml:"model"`
|
|
Streaming bool `yaml:"streaming"`
|
|
ClientType string `yaml:"client"`
|
|
}
|
|
|
|
type PromptTemplate struct {
|
|
TargetTokens int `yaml:"target_tokens"`
|
|
Templates []string `yaml:"templates"`
|
|
}
|
|
|
|
type RequestMix struct {
|
|
Type string `yaml:"type"`
|
|
Weight float64 `yaml:"weight"`
|
|
}
|
|
|
|
type ConcurrencyConfig struct {
|
|
Steps []int `yaml:"steps"`
|
|
DurationPerStep int `yaml:"duration_per_step"`
|
|
MaxGoroutines int `yaml:"max_goroutines"`
|
|
}
|
|
|
|
type Config struct {
|
|
API APIConfig `yaml:"api"`
|
|
PromptTemplates map[string]PromptTemplate `yaml:"prompt_templates"`
|
|
Requests []RequestMix `yaml:"requests"`
|
|
Concurrency ConcurrencyConfig `yaml:"concurrency"`
|
|
Timeout int `yaml:"timeout"`
|
|
PoissonLambda float64 `yaml:"poisson_lambda"`
|
|
TokenizerModel string `yaml:"tokenizer_model"`
|
|
}
|
|
|
|
func LoadConfig(path string) (*Config, error) {
|
|
data, err := ioutil.ReadFile(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var cfg Config
|
|
if err := yaml.Unmarshal(data, &cfg); err != nil {
|
|
return nil, err
|
|
}
|
|
return &cfg, nil
|
|
}
|