redis缓存替换+pgvector向量替换

ai-chat-service/pkg/config/config.go

@@ -1,7 +1,11 @@
 package config
 
 import (
+	"bufio"
 	"log"
+	"os"
+	"path/filepath"
+	"strings"
 
 	"github.com/spf13/viper"
 )
@@ -54,6 +58,35 @@ type Config struct {
 			Address string
 		}
 	}
+	Vector struct {
+		Provider  string
+		Threshold float32
+		Tencent   struct {
+			Url                string
+			Username           string
+			Pwd                string
+			Database           string
+			Timeout            int
+			MaxIdleConnPerHost int
+			ReadConsistency    string
+			IdleConnTimeout    int
+		}
+		Pgvector struct {
+			DSN         string `mapstructure:"dsn"`
+			Table       string `mapstructure:"table"`
+			Dimensions  int    `mapstructure:"dimensions"`
+			MaxLifeTime int    `mapstructure:"maxLifeTime"`
+			MaxOpenConn int    `mapstructure:"maxOpenConn"`
+			MaxIdleConn int    `mapstructure:"maxIdleConn"`
+		}
+	}
+	Embedding struct {
+		Provider string
+		BaseUrl  string `mapstructure:"base_url"`
+		ApiKey   string `mapstructure:"api_key"`
+		Model    string `mapstructure:"model"`
+		Timeout  int
+	}
 	VectorDB struct {
 		Url                string
 		Username           string
@@ -69,6 +102,7 @@ type Config struct {
 var conf *Config
 
 func InitConfig(filePath string, typ ...string) {
+	loadProjectDotEnv(filePath)
 	v := viper.New()
 	v.SetConfigFile(filePath)
 	if len(typ) > 0 {
@@ -83,9 +117,137 @@ func InitConfig(filePath string, typ ...string) {
 	if err != nil {
 		log.Fatal(err)
 	}
+	normalizeConfig(conf)
 
 }
 
 func GetConfig() *Config {
 	return conf
 }
+
+func normalizeConfig(conf *Config) {
+	if conf.Vector.Provider == "" {
+		conf.Vector.Provider = "tencent"
+	}
+	if conf.Vector.Threshold == 0 {
+		conf.Vector.Threshold = 0.99
+	}
+
+	// Backfill the new vector.tencent block from the legacy vectorDB config.
+	if conf.Vector.Tencent.Url == "" {
+		conf.Vector.Tencent.Url = conf.VectorDB.Url
+	}
+	if conf.Vector.Tencent.Username == "" {
+		conf.Vector.Tencent.Username = conf.VectorDB.Username
+	}
+	if conf.Vector.Tencent.Pwd == "" {
+		conf.Vector.Tencent.Pwd = conf.VectorDB.Pwd
+	}
+	if conf.Vector.Tencent.Database == "" {
+		conf.Vector.Tencent.Database = conf.VectorDB.Database
+	}
+	if conf.Vector.Tencent.Timeout == 0 {
+		conf.Vector.Tencent.Timeout = conf.VectorDB.Timeout
+	}
+	if conf.Vector.Tencent.MaxIdleConnPerHost == 0 {
+		conf.Vector.Tencent.MaxIdleConnPerHost = conf.VectorDB.MaxIdleConnPerHost
+	}
+	if conf.Vector.Tencent.ReadConsistency == "" {
+		conf.Vector.Tencent.ReadConsistency = conf.VectorDB.ReadConsistency
+	}
+	if conf.Vector.Tencent.IdleConnTimeout == 0 {
+		conf.Vector.Tencent.IdleConnTimeout = conf.VectorDB.IdleConnTimeout
+	}
+
+	if conf.Embedding.Provider == "" {
+		conf.Embedding.Provider = "openai-compatible"
+	}
+	if conf.Embedding.BaseUrl == "" {
+		conf.Embedding.BaseUrl = conf.Chat.BaseUrl
+	}
+	if conf.Embedding.ApiKey == "" {
+		conf.Embedding.ApiKey = conf.Chat.ApiKey
+	}
+	if conf.Embedding.Timeout == 0 {
+		conf.Embedding.Timeout = 10
+	}
+	overrideChatFromEnv(conf)
+	overrideEmbeddingFromEnv(conf)
+}
+
+func overrideChatFromEnv(conf *Config) {
+	if value := os.Getenv("AI_CHAT_OPENAI_BASE_URL"); value != "" {
+		conf.Chat.BaseUrl = value
+	} else if value := os.Getenv("OPENAI_BASE_URL"); value != "" {
+		conf.Chat.BaseUrl = value
+	}
+
+	if value := os.Getenv("AI_CHAT_OPENAI_MODEL"); value != "" {
+		conf.Chat.Model = value
+	} else if value := os.Getenv("OPENAI_MODEL"); value != "" {
+		conf.Chat.Model = value
+	}
+
+	if value := os.Getenv("AI_CHAT_OPENAI_API_KEY"); value != "" {
+		conf.Chat.ApiKey = value
+		return
+	}
+	if value := os.Getenv("OPENAI_API_KEY"); value != "" {
+		conf.Chat.ApiKey = value
+		return
+	}
+	if value := os.Getenv("MOONSHOT_API_KEY"); value != "" {
+		conf.Chat.ApiKey = value
+	}
+}
+
+func overrideEmbeddingFromEnv(conf *Config) {
+	if value := os.Getenv("AI_CHAT_EMBEDDING_BASE_URL"); value != "" {
+		conf.Embedding.BaseUrl = value
+	}
+	if value := os.Getenv("AI_CHAT_EMBEDDING_MODEL"); value != "" {
+		conf.Embedding.Model = value
+	}
+	if value := os.Getenv("AI_CHAT_EMBEDDING_API_KEY"); value != "" {
+		conf.Embedding.ApiKey = value
+		return
+	}
+	if value := os.Getenv("ZAI_API_KEY"); value != "" {
+		conf.Embedding.ApiKey = value
+	}
+}
+
+func loadProjectDotEnv(configFilePath string) {
+	projectRoot := filepath.Dir(filepath.Dir(configFilePath))
+	loadDotEnvFile(filepath.Join(projectRoot, ".env"))
+}
+
+func loadDotEnvFile(path string) {
+	file, err := os.Open(path)
+	if err != nil {
+		return
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		if line == "" || strings.HasPrefix(line, "#") {
+			continue
+		}
+		key, value, ok := strings.Cut(line, "=")
+		if !ok {
+			continue
+		}
+		key = strings.TrimSpace(key)
+		value = strings.TrimSpace(value)
+		value = strings.Trim(value, `"'`)
+		if key == "" {
+			continue
+		}
+		if _, exists := os.LookupEnv(key); exists {
+			continue
+		}
+		_ = os.Setenv(key, value)
+	}
+}