ai_old/agent.go
2024-09-07 23:17:32 +08:00

90 lines
2.0 KiB
Go

package ai
import (
"apigo.cc/ai/agent"
"github.com/ssgo/config"
"github.com/ssgo/u"
"sync"
)
var confAes = u.NewAes([]byte("?GQ$0K0GgLdO=f+~L68PLm$uhKr4'=tV"), []byte("VFs7@sK61cj^f?HZ"))
var keysIsSet = false
func SetSSKey(key, iv []byte) {
if !keysIsSet {
confAes = u.NewAes(key, iv)
keysIsSet = true
}
}
const (
TypeText = agent.TypeText
TypeImage = agent.TypeImage
TypeVideo = agent.TypeVideo
RoleSystem = agent.RoleSystem
RoleUser = agent.RoleUser
RoleAssistant = agent.RoleAssistant
RoleTool = agent.RoleTool
ToolCodeInterpreter = agent.ToolCodeInterpreter
ToolWebSearch = agent.ToolWebSearch
)
type APIConfig struct {
Endpoint string
ApiKey string
DefaultChatModelConfig ChatModelConfig
}
type AgentConfig struct {
ApiKey string
Endpoint string
Agent string
ChatConfig ChatModelConfig
}
var agentConfigs map[string]*AgentConfig
var agentConfigsLock = sync.RWMutex{}
func GetAgent(name string) agent.Agent {
ag := agent.GetAgent(name)
if ag != nil {
return ag
}
var agConf *AgentConfig
if agentConfigs == nil {
agConfs := make(map[string]*AgentConfig)
config.LoadConfig("agent", &agConfs)
agConf = agConfs[name]
if agConf != nil {
agentConfigsLock.Lock()
agentConfigs = agConfs
agentConfigsLock.Unlock()
}
} else {
agentConfigsLock.RLock()
agConf = agentConfigs[name]
agentConfigsLock.RUnlock()
}
if agConf == nil {
return nil
}
if agConf.Agent == "" {
agConf.Agent = name
}
return agent.CreateAgent(name, agConf.Agent, agent.APIConfig{
Endpoint: agConf.Endpoint,
ApiKey: confAes.DecryptUrlBase64ToString(agConf.ApiKey),
DefaultChatModelConfig: agent.ChatModelConfig{
Model: agConf.ChatConfig.Model,
MaxTokens: agConf.ChatConfig.MaxTokens,
Temperature: agConf.ChatConfig.Temperature,
TopP: agConf.ChatConfig.TopP,
Tools: agConf.ChatConfig.Tools,
},
})
}