132 lines
3.2 KiB
Go
132 lines
3.2 KiB
Go
package llm
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"aichat/config"
|
|
|
|
ark "github.com/volcengine/volcengine-go-sdk/service/arkruntime"
|
|
)
|
|
|
|
type Profile struct {
|
|
Config config.OpenAIConfig
|
|
Client *ark.Client
|
|
}
|
|
|
|
type State struct {
|
|
mu sync.RWMutex
|
|
profiles map[string]*Profile
|
|
order []string
|
|
activeName string
|
|
}
|
|
|
|
type ListResponse struct {
|
|
Active string `json:"active"`
|
|
Profiles []config.OpenAIConfig `json:"profiles"`
|
|
}
|
|
|
|
func NewState(configs []config.OpenAIConfig) (*State, error) {
|
|
state := &State{
|
|
profiles: make(map[string]*Profile, len(configs)),
|
|
order: make([]string, 0, len(configs)),
|
|
}
|
|
for _, item := range configs {
|
|
if strings.TrimSpace(item.Name) == "" {
|
|
return nil, errors.New("openai.name 不能为空")
|
|
}
|
|
if strings.TrimSpace(item.APIKey) == "" {
|
|
return nil, fmt.Errorf("openai.%s.api_key 未配置，也未设置环境变量 ARK_API_KEY", item.Name)
|
|
}
|
|
if strings.TrimSpace(item.Model) == "" {
|
|
return nil, fmt.Errorf("openai.%s.model 未配置", item.Name)
|
|
}
|
|
if strings.TrimSpace(item.BaseURL) == "" {
|
|
return nil, fmt.Errorf("openai.%s.base_url 未配置", item.Name)
|
|
}
|
|
if item.Timeout <= 0 {
|
|
return nil, fmt.Errorf("openai.%s.timeout 必须大于 0", item.Name)
|
|
}
|
|
if _, ok := state.profiles[item.Name]; ok {
|
|
return nil, fmt.Errorf("openai 配置名称重复: %s", item.Name)
|
|
}
|
|
state.profiles[item.Name] = &Profile{
|
|
Config: item,
|
|
Client: ark.NewClientWithApiKey(
|
|
item.APIKey,
|
|
ark.WithBaseUrl(item.BaseURL),
|
|
ark.WithTimeout(time.Duration(item.Timeout)*time.Second),
|
|
),
|
|
}
|
|
state.order = append(state.order, item.Name)
|
|
if item.Active && state.activeName == "" {
|
|
state.activeName = item.Name
|
|
}
|
|
}
|
|
if len(state.order) == 0 {
|
|
return nil, errors.New("openai 配置不能为空")
|
|
}
|
|
if state.activeName == "" {
|
|
state.activeName = state.order[0]
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func (s *State) ActiveProfile() *Profile {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
return s.profiles[s.activeName]
|
|
}
|
|
|
|
func (s *State) GetProfile(name string) (*Profile, error) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
if strings.TrimSpace(name) == "" {
|
|
return s.profiles[s.activeName], nil
|
|
}
|
|
profile, ok := s.profiles[strings.TrimSpace(name)]
|
|
if !ok {
|
|
return nil, fmt.Errorf("OpenAI 配置不存在: %s", name)
|
|
}
|
|
return profile, nil
|
|
}
|
|
|
|
func (s *State) SwitchActive(name string) (*Profile, error) {
|
|
name = strings.TrimSpace(name)
|
|
if name == "" {
|
|
return nil, errors.New("OpenAI 配置名称不能为空")
|
|
}
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
profile, ok := s.profiles[name]
|
|
if !ok {
|
|
return nil, fmt.Errorf("OpenAI 配置不存在: %s", name)
|
|
}
|
|
s.activeName = name
|
|
return profile, nil
|
|
}
|
|
|
|
func (s *State) ListProfiles() ListResponse {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
profiles := make([]config.OpenAIConfig, 0, len(s.order))
|
|
for _, name := range s.order {
|
|
profile := s.profiles[name]
|
|
cfg := profile.Config
|
|
cfg.APIKey = ""
|
|
cfg.Active = name == s.activeName
|
|
profiles = append(profiles, cfg)
|
|
}
|
|
return ListResponse{Active: s.activeName, Profiles: profiles}
|
|
}
|
|
|
|
func PublicConfig(profile *Profile, active bool) config.OpenAIConfig {
|
|
cfg := profile.Config
|
|
cfg.APIKey = ""
|
|
cfg.Active = active
|
|
return cfg
|
|
}
|