factory.go 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081
  1. package embedding
  2. import (
  3. "fmt"
  4. "os"
  5. )
  6. // EmbeddingFactory 嵌入服务工厂
  7. type EmbeddingFactory struct{}
  8. // NewEmbeddingFactory 创建嵌入服务工厂
  9. func NewEmbeddingFactory() *EmbeddingFactory {
  10. return &EmbeddingFactory{}
  11. }
  12. // getConfigFromEnv 从环境变量读取配置
  13. func (f *EmbeddingFactory) getConfigFromEnv() (embeddingType, apiURL, apiKey, model string) {
  14. embeddingType = os.Getenv("EMBEDDING_TYPE")
  15. if embeddingType == "" {
  16. embeddingType = "local" // 默认使用本地 BGE
  17. }
  18. apiURL = os.Getenv("EMBEDDING_API_URL")
  19. apiKey = os.Getenv("EMBEDDING_API_KEY")
  20. model = os.Getenv("EMBEDDING_MODEL")
  21. // BGE 配置
  22. if embeddingType == "local" || embeddingType == "bge" {
  23. if bgeURL := os.Getenv("BGE_API_URL"); bgeURL != "" {
  24. apiURL = bgeURL
  25. }
  26. if bgeKey := os.Getenv("BGE_API_KEY"); bgeKey != "" {
  27. apiKey = bgeKey
  28. }
  29. if bgeModel := os.Getenv("BGE_MODEL_NAME"); bgeModel != "" {
  30. model = bgeModel
  31. }
  32. }
  33. return
  34. }
  35. // CreateDefaultEmbeddingService 创建默认嵌入服务
  36. // 优先尝试 BGE 本地服务,如果失败则使用 OpenAI
  37. func (f *EmbeddingFactory) CreateDefaultEmbeddingService() (EmbeddingService, error) {
  38. embeddingType, apiURL, apiKey, model := f.getConfigFromEnv()
  39. switch embeddingType {
  40. case "openai", "api":
  41. if apiKey == "" {
  42. return nil, fmt.Errorf("EMBEDDING_API_KEY 未设置")
  43. }
  44. return NewOpenAIEmbeddingService(apiURL, apiKey, model), nil
  45. case "local", "bge":
  46. // 尝试创建 BGE 服务
  47. return NewBGEEmbeddingService(apiURL, apiKey, model), nil
  48. default:
  49. return nil, fmt.Errorf("不支持的嵌入服务类型: %s", embeddingType)
  50. }
  51. }
  52. // CreateEmbeddingService 根据类型创建嵌入服务
  53. func (f *EmbeddingFactory) CreateEmbeddingService(embeddingType string) (EmbeddingService, error) {
  54. _, apiURL, apiKey, model := f.getConfigFromEnv()
  55. switch embeddingType {
  56. case "openai", "api":
  57. if apiKey == "" {
  58. return nil, fmt.Errorf("EMBEDDING_API_KEY 未设置")
  59. }
  60. return NewOpenAIEmbeddingService(apiURL, apiKey, model), nil
  61. case "local", "bge":
  62. return NewBGEEmbeddingService(apiURL, apiKey, model), nil
  63. default:
  64. return nil, fmt.Errorf("不支持的嵌入服务类型: %s", embeddingType)
  65. }
  66. }