146 lines
3.2 KiB
Go
146 lines
3.2 KiB
Go
package llm
|
|
|
|
type ChatMessage struct {
|
|
Role string
|
|
Contents []ChatMessageContent
|
|
}
|
|
|
|
type ChatMessageContent struct {
|
|
Type string // text, image, audio, video
|
|
Content string
|
|
}
|
|
|
|
type ChatConfig struct {
|
|
defaultConfig *ChatConfig
|
|
Model string
|
|
Ratio float64
|
|
MaxTokens int
|
|
Temperature float64
|
|
TopP float64
|
|
Tools map[string]any
|
|
}
|
|
|
|
func (chatConfig *ChatConfig) SetDefault(config *ChatConfig) {
|
|
chatConfig.defaultConfig = config
|
|
}
|
|
|
|
func (chatConfig *ChatConfig) GetModel() string {
|
|
if chatConfig.Model == "" && chatConfig.defaultConfig != nil {
|
|
return chatConfig.defaultConfig.Model
|
|
}
|
|
return chatConfig.Model
|
|
}
|
|
|
|
func (chatConfig *ChatConfig) GetMaxTokens() int {
|
|
if chatConfig.MaxTokens == 0 && chatConfig.defaultConfig != nil {
|
|
return chatConfig.defaultConfig.MaxTokens
|
|
}
|
|
return chatConfig.MaxTokens
|
|
}
|
|
|
|
func (chatConfig *ChatConfig) GetTemperature() float64 {
|
|
if chatConfig.Temperature == 0 && chatConfig.defaultConfig != nil {
|
|
return chatConfig.defaultConfig.Temperature
|
|
}
|
|
return chatConfig.Temperature
|
|
}
|
|
|
|
func (chatConfig *ChatConfig) GetTopP() float64 {
|
|
if chatConfig.TopP == 0 && chatConfig.defaultConfig != nil {
|
|
return chatConfig.defaultConfig.TopP
|
|
}
|
|
return chatConfig.TopP
|
|
}
|
|
|
|
func (chatConfig *ChatConfig) GetTools() map[string]any {
|
|
if chatConfig.Tools == nil && chatConfig.defaultConfig != nil {
|
|
return chatConfig.defaultConfig.Tools
|
|
}
|
|
return chatConfig.Tools
|
|
}
|
|
|
|
type TokenUsage struct {
|
|
AskTokens int64
|
|
AnswerTokens int64
|
|
TotalTokens int64
|
|
}
|
|
|
|
type MessagesMaker struct {
|
|
list []ChatMessage
|
|
}
|
|
|
|
func Messages() *MessagesMaker {
|
|
return &MessagesMaker{
|
|
list: make([]ChatMessage, 0),
|
|
}
|
|
}
|
|
|
|
func (m *MessagesMaker) Make() []ChatMessage {
|
|
return m.list
|
|
}
|
|
|
|
func (m *MessagesMaker) User(contents ...ChatMessageContent) *MessagesMaker {
|
|
m.list = append(m.list, ChatMessage{
|
|
Role: RoleUser,
|
|
Contents: contents,
|
|
})
|
|
return m
|
|
}
|
|
|
|
func (m *MessagesMaker) Assistant(contents ...ChatMessageContent) *MessagesMaker {
|
|
m.list = append(m.list, ChatMessage{
|
|
Role: RoleAssistant,
|
|
Contents: contents,
|
|
})
|
|
return m
|
|
}
|
|
|
|
func (m *MessagesMaker) System(contents ...ChatMessageContent) *MessagesMaker {
|
|
m.list = append(m.list, ChatMessage{
|
|
Role: RoleSystem,
|
|
Contents: contents,
|
|
})
|
|
return m
|
|
}
|
|
|
|
func (m *MessagesMaker) Tool(contents ...ChatMessageContent) *MessagesMaker {
|
|
m.list = append(m.list, ChatMessage{
|
|
Role: RoleTool,
|
|
Contents: contents,
|
|
})
|
|
return m
|
|
}
|
|
|
|
func (m *MessagesMaker) Text(text string) *MessagesMaker {
|
|
if len(m.list) > 0 {
|
|
lastIndex := len(m.list) - 1
|
|
m.list[lastIndex].Contents = append(m.list[lastIndex].Contents, ChatMessageContent{
|
|
Type: TypeText,
|
|
Content: text,
|
|
})
|
|
}
|
|
return m
|
|
}
|
|
|
|
func (m *MessagesMaker) Image(text string) *MessagesMaker {
|
|
if len(m.list) > 0 {
|
|
lastIndex := len(m.list) - 1
|
|
m.list[lastIndex].Contents = append(m.list[lastIndex].Contents, ChatMessageContent{
|
|
Type: TypeText,
|
|
Content: text,
|
|
})
|
|
}
|
|
return m
|
|
}
|
|
|
|
func (m *MessagesMaker) Video(text string) *MessagesMaker {
|
|
if len(m.list) > 0 {
|
|
lastIndex := len(m.list) - 1
|
|
m.list[lastIndex].Contents = append(m.list[lastIndex].Contents, ChatMessageContent{
|
|
Type: TypeText,
|
|
Content: text,
|
|
})
|
|
}
|
|
return m
|
|
}
|