162 lines
5.0 KiB
Go
162 lines
5.0 KiB
Go
package config
|
|
|
|
import (
|
|
"os"
|
|
"strings"
|
|
|
|
"github.com/spf13/pflag"
|
|
"gopkg.in/yaml.v3"
|
|
)
|
|
|
|
type AppConfigFile struct {
|
|
Server struct {
|
|
Port string `yaml:"port"`
|
|
LogLevel string `yaml:"log_level"`
|
|
} `yaml:"server"`
|
|
AITrigger struct {
|
|
AnalysisThreshold float64 `yaml:"analysis_threshold"`
|
|
} `yaml:"ai_trigger"`
|
|
Prompts struct {
|
|
SystemPromptPath string `yaml:"system_prompt_path"`
|
|
UserPromptPath string `yaml:"user_prompt_path"`
|
|
} `yaml:"prompts"`
|
|
}
|
|
|
|
type Config struct {
|
|
Port string
|
|
DBPath string
|
|
TargetConfig string
|
|
FilterConfig string
|
|
AppConfig string
|
|
LogLevel string
|
|
LLMProvider string // "gemini", "deepseek", "mock"
|
|
GeminiAPIKey string
|
|
DeepSeekAPIKey string
|
|
DeepSeekBaseURL string
|
|
AnalysisThreshold float64
|
|
SystemPromptPath string
|
|
UserPromptPath string
|
|
SystemPrompt string
|
|
UserPrompt string
|
|
}
|
|
|
|
func init() {
|
|
LoadEnvFile(".env")
|
|
}
|
|
|
|
// LoadEnvFile parses a simple KEY=VALUE .env file if it exists without overriding existing env vars
|
|
func LoadEnvFile(filenames ...string) {
|
|
if len(filenames) == 0 {
|
|
filenames = []string{".env"}
|
|
}
|
|
for _, filename := range filenames {
|
|
data, err := os.ReadFile(filename)
|
|
if err != nil {
|
|
continue
|
|
}
|
|
lines := strings.Split(string(data), "\n")
|
|
for _, line := range lines {
|
|
line = strings.TrimSpace(line)
|
|
if line == "" || strings.HasPrefix(line, "#") {
|
|
continue
|
|
}
|
|
parts := strings.SplitN(line, "=", 2)
|
|
if len(parts) == 2 {
|
|
key := strings.TrimSpace(parts[0])
|
|
val := strings.TrimSpace(parts[1])
|
|
val = strings.Trim(val, `"'`)
|
|
if os.Getenv(key) == "" {
|
|
os.Setenv(key, val)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func RegisterFlags(flags *pflag.FlagSet) *Config {
|
|
cfg := &Config{}
|
|
|
|
flags.StringVarP(&cfg.Port, "port", "p", getEnvOrDefault("PORT", "8080"), "HTTP server port")
|
|
flags.StringVarP(&cfg.DBPath, "db", "d", getEnvOrDefault("DB_PATH", "data/db/hireme.db"), "Path to database file")
|
|
flags.StringVarP(&cfg.TargetConfig, "target-config", "c", getEnvOrDefault("TARGET_CONFIG", "data/config/targets.yaml"), "Path to targets YAML file")
|
|
flags.StringVarP(&cfg.FilterConfig, "filter-config", "f", getEnvOrDefault("FILTER_CONFIG", "data/config/filters.yaml"), "Path to filters YAML file")
|
|
flags.StringVar(&cfg.AppConfig, "app-config", getEnvOrDefault("APP_CONFIG", "data/config/config.yaml"), "Path to app YAML config file")
|
|
flags.StringVar(&cfg.LogLevel, "log-level", getEnvOrDefault("LOG_LEVEL", "info"), "Logging level (trace, debug, info, error)")
|
|
flags.StringVarP(&cfg.LLMProvider, "llm-provider", "l", getEnvOrDefault("LLM_PROVIDER", ""), "Active LLM Provider (gemini, deepseek, mock)")
|
|
flags.StringVar(&cfg.GeminiAPIKey, "gemini-api-key", os.Getenv("GEMINI_API_KEY"), "Gemini API Key")
|
|
flags.StringVar(&cfg.DeepSeekAPIKey, "deepseek-api-key", os.Getenv("DEEPSEEK_API_KEY"), "DeepSeek API Key")
|
|
flags.StringVar(&cfg.DeepSeekBaseURL, "deepseek-base-url", getEnvOrDefault("DEEPSEEK_BASE_URL", "https://api.deepseek.com"), "DeepSeek Base URL")
|
|
flags.Float64Var(&cfg.AnalysisThreshold, "analysis-threshold", 25.0, "Threshold score from filter to trigger LLM analysis")
|
|
|
|
return cfg
|
|
}
|
|
|
|
func (c *Config) Normalize() {
|
|
// Re-check env vars in case .env was loaded after initial flag binding
|
|
if c.DeepSeekAPIKey == "" {
|
|
c.DeepSeekAPIKey = os.Getenv("DEEPSEEK_API_KEY")
|
|
}
|
|
if c.GeminiAPIKey == "" {
|
|
c.GeminiAPIKey = os.Getenv("GEMINI_API_KEY")
|
|
}
|
|
|
|
// Load data/config/config.yaml if present
|
|
if c.AppConfig != "" {
|
|
if data, err := os.ReadFile(c.AppConfig); err == nil {
|
|
var appCfg AppConfigFile
|
|
if err := yaml.Unmarshal(data, &appCfg); err == nil {
|
|
if appCfg.Server.Port != "" && c.Port == "8080" {
|
|
c.Port = appCfg.Server.Port
|
|
}
|
|
if appCfg.Server.LogLevel != "" && c.LogLevel == "info" {
|
|
c.LogLevel = appCfg.Server.LogLevel
|
|
}
|
|
if appCfg.AITrigger.AnalysisThreshold > 0 {
|
|
c.AnalysisThreshold = appCfg.AITrigger.AnalysisThreshold
|
|
}
|
|
if appCfg.Prompts.SystemPromptPath != "" {
|
|
c.SystemPromptPath = appCfg.Prompts.SystemPromptPath
|
|
}
|
|
if appCfg.Prompts.UserPromptPath != "" {
|
|
c.UserPromptPath = appCfg.Prompts.UserPromptPath
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if c.SystemPromptPath == "" {
|
|
c.SystemPromptPath = "data/config/prompts/system_prompt.txt"
|
|
}
|
|
if c.UserPromptPath == "" {
|
|
c.UserPromptPath = "data/config/prompts/user_prompt.txt"
|
|
}
|
|
|
|
// Load prompt contents from files
|
|
if promptData, err := os.ReadFile(c.SystemPromptPath); err == nil {
|
|
c.SystemPrompt = strings.TrimSpace(string(promptData))
|
|
}
|
|
if userPromptData, err := os.ReadFile(c.UserPromptPath); err == nil {
|
|
c.UserPrompt = strings.TrimSpace(string(userPromptData))
|
|
}
|
|
|
|
// Determine LLM provider
|
|
if c.LLMProvider == "" {
|
|
if c.DeepSeekAPIKey != "" {
|
|
c.LLMProvider = "deepseek"
|
|
} else if c.GeminiAPIKey != "" {
|
|
c.LLMProvider = "gemini"
|
|
} else {
|
|
c.LLMProvider = "mock"
|
|
}
|
|
}
|
|
c.LLMProvider = strings.ToLower(c.LLMProvider)
|
|
}
|
|
|
|
func getEnvOrDefault(key, defaultValue string) string {
|
|
val := os.Getenv(key)
|
|
if val != "" {
|
|
return val
|
|
}
|
|
return defaultValue
|
|
}
|