97 lines
2.7 KiB
Go
97 lines
2.7 KiB
Go
|
package llm
|
||
|
|
|
||
|
|
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 Config struct {
|
||
|
|
Endpoint string
|
||
|
|
ApiKey string
|
||
|
|
ChatConfig ChatConfig
|
||
|
|
GCConfig GCConfig
|
||
|
|
}
|
||
|
|
|
||
|
|
type LLM interface {
|
||
|
|
Support() Support
|
||
|
|
Ask(messages []ChatMessage, config ChatConfig, 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(prompt string, config GCConfig) ([]string, error)
|
||
|
|
FastMakeImage(prompt string, config GCConfig) ([]string, error)
|
||
|
|
BestMakeImage(prompt string, config GCConfig) ([]string, error)
|
||
|
|
MakeVideo(prompt string, config GCConfig) ([]string, []string, error)
|
||
|
|
FastMakeVideo(prompt string, config GCConfig) ([]string, []string, error)
|
||
|
|
BestMakeVideo(prompt string, config GCConfig) ([]string, []string, error)
|
||
|
|
}
|
||
|
|
|
||
|
|
var llmMakers = map[string]func(Config) LLM{}
|
||
|
|
var llmMakersLock = sync.RWMutex{}
|
||
|
|
|
||
|
|
var llms = map[string]LLM{}
|
||
|
|
var llmsLock = sync.RWMutex{}
|
||
|
|
|
||
|
|
func Register(llmId string, maker func(Config) LLM) {
|
||
|
|
llmMakersLock.Lock()
|
||
|
|
llmMakers[llmId] = maker
|
||
|
|
llmMakersLock.Unlock()
|
||
|
|
}
|
||
|
|
|
||
|
|
func Create(name, llmId string, config Config) LLM {
|
||
|
|
llmMakersLock.RLock()
|
||
|
|
maker := llmMakers[llmId]
|
||
|
|
llmMakersLock.RUnlock()
|
||
|
|
|
||
|
|
if maker != nil {
|
||
|
|
llm := maker(config)
|
||
|
|
llmsLock.Lock()
|
||
|
|
llms[name] = llm
|
||
|
|
llmsLock.Unlock()
|
||
|
|
return llm
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func Get(name string) LLM {
|
||
|
|
llmsLock.RLock()
|
||
|
|
llm := llms[name]
|
||
|
|
llmsLock.RUnlock()
|
||
|
|
return llm
|
||
|
|
}
|
||
|
|
|
||
|
|
func List() map[string]LLM {
|
||
|
|
list := map[string]LLM{}
|
||
|
|
llmsLock.RLock()
|
||
|
|
for name, llm := range llms {
|
||
|
|
list[name] = llm
|
||
|
|
}
|
||
|
|
llmsLock.RUnlock()
|
||
|
|
return list
|
||
|
|
}
|