ai_scheduler/internal/config/config.go

110 lines
3.8 KiB
Go

package config
import (
"fmt"
"time"
"github.com/spf13/viper"
)
// Config 应用配置
type Config struct {
Server ServerConfig `mapstructure:"server"`
Ollama OllamaConfig `mapstructure:"ollama"`
Sys SysConfig `mapstructure:"sys"`
Tools ToolsConfig `mapstructure:"tools"`
Logging LoggingConfig `mapstructure:"logging"`
Redis *Redis `protobuf:"bytes,1,opt,name=Redis,proto3" json:"Redis,omitempty"`
DB *DB `protobuf:"bytes,3,opt,name=TransDB,proto3" json:"TransDB,omitempty"`
}
// SysConfig 系统配置
type SysConfig struct {
SessionLen int `mapstructure:"session_len"`
}
// ServerConfig 服务器配置
type ServerConfig struct {
Port int `mapstructure:"port"`
Host string `mapstructure:"host"`
}
// OllamaConfig Ollama配置
type OllamaConfig struct {
BaseURL string `mapstructure:"base_url"`
Model string `mapstructure:"model"`
Timeout time.Duration `mapstructure:"timeout"`
}
type Redis struct {
Host string `protobuf:"bytes,1,opt,name=host,proto3" json:"host,omitempty"`
Type string `protobuf:"bytes,2,opt,name=type,proto3" json:"type,omitempty"`
Pass string `protobuf:"bytes,3,opt,name=pass,proto3" json:"pass,omitempty"`
Key string `protobuf:"bytes,4,opt,name=key,proto3" json:"key,omitempty"`
Tls int32 `protobuf:"varint,5,opt,name=tls,proto3" json:"tls,omitempty"`
Db int32 `protobuf:"varint,6,opt,name=db,proto3" json:"db,omitempty"`
MaxIdle int32 `protobuf:"varint,7,opt,name=maxIdle,proto3" json:"maxIdle,omitempty"`
PoolSize int32 `protobuf:"varint,8,opt,name=poolSize,proto3" json:"poolSize,omitempty"`
MaxIdleTime int32 `protobuf:"varint,9,opt,name=maxIdleTime,proto3" json:"maxIdleTime,omitempty"`
}
type DB struct {
Driver string `protobuf:"bytes,1,opt,name=driver,proto3" json:"driver,omitempty"`
Source string `protobuf:"bytes,2,opt,name=source,proto3" json:"source,omitempty"`
MaxIdle int32 `protobuf:"varint,3,opt,name=maxIdle,proto3" json:"maxIdle,omitempty"`
MaxOpen int32 `protobuf:"varint,4,opt,name=maxOpen,proto3" json:"maxOpen,omitempty"`
MaxLifetime int32 `protobuf:"varint,5,opt,name=maxLifetime,proto3" json:"maxLifetime,omitempty"`
IsDebug bool `protobuf:"varint,6,opt,name=isDebug,proto3" json:"isDebug,omitempty"`
}
// ToolsConfig 工具配置
type ToolsConfig struct {
Weather ToolConfig `mapstructure:"weather"`
Calculator ToolConfig `mapstructure:"calculator"`
ZltxOrderDetail ToolConfig `mapstructure:"zltxOrderDetail"`
ZltxOrderLog ToolConfig `mapstructure:"zltxOrderLog"`
Knowledge ToolConfig `mapstructure:"knowledge"`
}
// ToolConfig 单个工具配置
type ToolConfig struct {
Enabled bool `mapstructure:"enabled"`
BaseURL string `mapstructure:"base_url"`
APIKey string `mapstructure:"api_key"`
BizSystem string `mapstructure:"biz_system"`
}
// LoggingConfig 日志配置
type LoggingConfig struct {
Level string `mapstructure:"level"`
Format string `mapstructure:"format"`
}
// LoadConfig 加载配置
func LoadConfig(configPath string) (*Config, error) {
viper.SetConfigFile(configPath)
viper.SetConfigType("yaml")
// 设置默认值
//viper.SetDefault("server.port", "8080")
//viper.SetDefault("server.host", "0.0.0.0")
//viper.SetDefault("ollama.base_url", "http://localhost:11434")
//viper.SetDefault("ollama.model", "llama2")
//viper.SetDefault("ollama.timeout", "30s")
//viper.SetDefault("logging.level", "info")
//viper.SetDefault("logging.format", "json")
// 读取配置文件
if err := viper.ReadInConfig(); err != nil {
return nil, fmt.Errorf("failed to read config file: %w", err)
}
// 解析配置
var config Config
if err := viper.Unmarshal(&config); err != nil {
return nil, fmt.Errorf("failed to unmarshal config: %w", err)
}
return &config, nil
}