92 lines
2.8 KiB
Go
92 lines
2.8 KiB
Go
|
package agent
|
||
|
|
|
||
|
|
import "sync"
|
||
|
|
|
||
|
|
const (
|
||
|
|
TypeText = "text"
|
||
|
|
TypeImage = "image"
|
||
|
|
TypeVideo = "video"
|
||
|
|
RoleSystem = "system"
|
||
|
|
RoleUser = "user"
|
||
|
|
RoleAssistant = "assistant"
|
||
|
|
RoleTool = "tool"
|
||
|
|
ToolCodeInterpreter = "codeInterpreter"
|
||
|
|
ToolWebSearch = "webSearch"
|
||
|
|
)
|
||
|
|
|
||
|
|
type Support struct {
|
||
|
|
Ask bool
|
||
|
|
AskWithImage bool
|
||
|
|
AskWithVideo bool
|
||
|
|
AskWithCodeInterpreter bool
|
||
|
|
AskWithWebSearch bool
|
||
|
|
MakeImage bool
|
||
|
|
MakeVideo bool
|
||
|
|
Models []string
|
||
|
|
}
|
||
|
|
|
||
|
|
type Agent interface {
|
||
|
|
Support() Support
|
||
|
|
Ask(messages []ChatMessage, config *ChatModelConfig, callback func(answer string)) (string, TokenUsage, error)
|
||
|
|
FastAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
|
||
|
|
LongAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
|
||
|
|
BatterAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
|
||
|
|
BestAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
|
||
|
|
MultiAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
|
||
|
|
BestMultiAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
|
||
|
|
CodeInterpreterAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
|
||
|
|
WebSearchAsk(messages []ChatMessage, callback func(answer string)) (string, TokenUsage, error)
|
||
|
|
MakeImage(model, prompt, size, refImage string) ([]string, error)
|
||
|
|
FastMakeImage(prompt, size, refImage string) ([]string, error)
|
||
|
|
BestMakeImage(prompt, size, refImage string) ([]string, error)
|
||
|
|
MakeVideo(model, prompt, size, refImage string) ([]string, []string, error)
|
||
|
|
FastMakeVideo(prompt, size, refImage string) ([]string, []string, error)
|
||
|
|
BestMakeVideo(prompt, size, refImage string) ([]string, []string, error)
|
||
|
|
}
|
||
|
|
|
||
|
|
type APIConfig struct {
|
||
|
|
Endpoint string
|
||
|
|
ApiKey string
|
||
|
|
DefaultChatModelConfig ChatModelConfig
|
||
|
|
}
|
||
|
|
|
||
|
|
type TokenUsage struct {
|
||
|
|
AskTokens int64
|
||
|
|
AnswerTokens int64
|
||
|
|
TotalTokens int64
|
||
|
|
}
|
||
|
|
|
||
|
|
var agentMakers = map[string]func(APIConfig) Agent{}
|
||
|
|
var agentMakersLock = sync.RWMutex{}
|
||
|
|
|
||
|
|
var agents = map[string]Agent{}
|
||
|
|
var agentsLock = sync.RWMutex{}
|
||
|
|
|
||
|
|
func RegisterAgentMaker(agentId string, maker func(APIConfig) Agent) {
|
||
|
|
agentMakersLock.Lock()
|
||
|
|
agentMakers[agentId] = maker
|
||
|
|
agentMakersLock.Unlock()
|
||
|
|
}
|
||
|
|
|
||
|
|
func CreateAgent(name, agentId string, config APIConfig) Agent {
|
||
|
|
agentMakersLock.RLock()
|
||
|
|
maker := agentMakers[agentId]
|
||
|
|
agentMakersLock.RUnlock()
|
||
|
|
|
||
|
|
if maker != nil {
|
||
|
|
agent := maker(config)
|
||
|
|
agentsLock.Lock()
|
||
|
|
agents[name] = agent
|
||
|
|
agentsLock.Unlock()
|
||
|
|
return agent
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func GetAgent(name string) Agent {
|
||
|
|
agentsLock.RLock()
|
||
|
|
agent := agents[name]
|
||
|
|
agentsLock.RUnlock()
|
||
|
|
return agent
|
||
|
|
}
|