gemini.go 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152
  1. package model_setting
  2. import (
  3. "one-api/setting/config"
  4. )
  5. // GeminiSettings 定义Gemini模型的配置
  6. type GeminiSettings struct {
  7. SafetySettings map[string]string `json:"safety_settings"`
  8. VersionSettings map[string]string `json:"version_settings"`
  9. }
  10. // 默认配置
  11. var defaultGeminiSettings = GeminiSettings{
  12. SafetySettings: map[string]string{
  13. "default": "OFF",
  14. "HARM_CATEGORY_CIVIC_INTEGRITY": "BLOCK_NONE",
  15. },
  16. VersionSettings: map[string]string{
  17. "default": "v1beta",
  18. "gemini-1.0-pro": "v1",
  19. },
  20. }
  21. // 全局实例
  22. var geminiSettings = defaultGeminiSettings
  23. func init() {
  24. // 注册到全局配置管理器
  25. config.GlobalConfig.Register("gemini", &geminiSettings)
  26. }
  27. // GetGeminiSettings 获取Gemini配置
  28. func GetGeminiSettings() *GeminiSettings {
  29. return &geminiSettings
  30. }
  31. // GetGeminiSafetySetting 获取安全设置
  32. func GetGeminiSafetySetting(key string) string {
  33. if value, ok := geminiSettings.SafetySettings[key]; ok {
  34. return value
  35. }
  36. return geminiSettings.SafetySettings["default"]
  37. }
  38. // GetGeminiVersionSetting 获取版本设置
  39. func GetGeminiVersionSetting(key string) string {
  40. if value, ok := geminiSettings.VersionSettings[key]; ok {
  41. return value
  42. }
  43. return geminiSettings.VersionSettings["default"]
  44. }