llm_old/llm/chat.go

188 lines
3.9 KiB
Go
Raw Permalink Normal View History

2024-10-02 14:09:54 +08:00
package llm
import (
"bytes"
"encoding/binary"
"math"
)
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 Usage struct {
AskTokens int64
AnswerTokens int64
TotalTokens int64
UsedTime 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
}
func bin2float64(in []byte) []float64 {
buf := bytes.NewBuffer(in)
out := make([]float64, len(in)/4)
for i := 0; i < len(out); i++ {
var f float32
_ = binary.Read(buf, binary.LittleEndian, &f)
out[i] = float64(f)
}
return out
}
func Similarity(buf1, buf2 []byte) float64 {
a := bin2float64(buf1)
b := bin2float64(buf2)
if len(a) != len(b) {
return 0
}
var dotProduct, magnitudeA, magnitudeB float64
for i := 0; i < len(a); i++ {
dotProduct += a[i] * b[i]
magnitudeA += a[i] * a[i]
magnitudeB += b[i] * b[i]
}
magnitudeA = math.Sqrt(magnitudeA)
magnitudeB = math.Sqrt(magnitudeB)
if magnitudeA == 0 || magnitudeB == 0 {
return 0
}
return dotProduct / (magnitudeA * magnitudeB)
}