ai_old/llm/chat.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
}