-
-
Notifications
You must be signed in to change notification settings - Fork 537
/
message_prompt_template.go
131 lines (106 loc) · 4.51 KB
/
message_prompt_template.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package prompts
import (
"fmt"
"github.com/tmc/langchaingo/llms"
)
// SystemMessagePromptTemplate is a message formatter that returns a system message.
type SystemMessagePromptTemplate struct {
Prompt PromptTemplate
}
var _ MessageFormatter = SystemMessagePromptTemplate{}
// FormatMessages formats the message with the values given.
func (p SystemMessagePromptTemplate) FormatMessages(values map[string]any) ([]llms.ChatMessage, error) {
text, err := p.Prompt.Format(values)
return []llms.ChatMessage{llms.SystemChatMessage{Content: text}}, err
}
// GetInputVariables returns the input variables the prompt expects.
func (p SystemMessagePromptTemplate) GetInputVariables() []string {
return p.Prompt.InputVariables
}
// NewSystemMessagePromptTemplate creates a new system message prompt template.
func NewSystemMessagePromptTemplate(template string, inputVariables []string) SystemMessagePromptTemplate {
return SystemMessagePromptTemplate{
Prompt: NewPromptTemplate(template, inputVariables),
}
}
// AIMessagePromptTemplate is a message formatter that returns an AI message.
type AIMessagePromptTemplate struct {
Prompt PromptTemplate
}
var _ MessageFormatter = AIMessagePromptTemplate{}
// FormatMessages formats the message with the values given.
func (p AIMessagePromptTemplate) FormatMessages(values map[string]any) ([]llms.ChatMessage, error) {
text, err := p.Prompt.Format(values)
return []llms.ChatMessage{llms.AIChatMessage{Content: text}}, err
}
// GetInputVariables returns the input variables the prompt expects.
func (p AIMessagePromptTemplate) GetInputVariables() []string {
return p.Prompt.InputVariables
}
// NewAIMessagePromptTemplate creates a new AI message prompt template.
func NewAIMessagePromptTemplate(template string, inputVariables []string) AIMessagePromptTemplate {
return AIMessagePromptTemplate{
Prompt: NewPromptTemplate(template, inputVariables),
}
}
// HumanMessagePromptTemplate is a message formatter that returns a human message.
type HumanMessagePromptTemplate struct {
Prompt PromptTemplate
}
var _ MessageFormatter = HumanMessagePromptTemplate{}
// FormatMessages formats the message with the values given.
func (p HumanMessagePromptTemplate) FormatMessages(values map[string]any) ([]llms.ChatMessage, error) {
text, err := p.Prompt.Format(values)
return []llms.ChatMessage{llms.HumanChatMessage{Content: text}}, err
}
// GetInputVariables returns the input variables the prompt expects.
func (p HumanMessagePromptTemplate) GetInputVariables() []string {
return p.Prompt.InputVariables
}
// NewHumanMessagePromptTemplate creates a new human message prompt template.
func NewHumanMessagePromptTemplate(template string, inputVariables []string) HumanMessagePromptTemplate {
return HumanMessagePromptTemplate{
Prompt: NewPromptTemplate(template, inputVariables),
}
}
// GenericMessagePromptTemplate is a message formatter that returns message with the specified speaker.
type GenericMessagePromptTemplate struct {
Prompt PromptTemplate
Role string
}
var _ MessageFormatter = GenericMessagePromptTemplate{}
// FormatMessages formats the message with the values given.
func (p GenericMessagePromptTemplate) FormatMessages(values map[string]any) ([]llms.ChatMessage, error) {
text, err := p.Prompt.Format(values)
return []llms.ChatMessage{llms.GenericChatMessage{Content: text, Role: p.Role}}, err
}
// GetInputVariables returns the input variables the prompt expects.
func (p GenericMessagePromptTemplate) GetInputVariables() []string {
return p.Prompt.InputVariables
}
// NewGenericMessagePromptTemplate creates a new generic message prompt template.
func NewGenericMessagePromptTemplate(role, template string, inputVariables []string) GenericMessagePromptTemplate {
return GenericMessagePromptTemplate{
Prompt: NewPromptTemplate(template, inputVariables),
Role: role,
}
}
type MessagesPlaceholder struct {
VariableName string
}
// FormatMessages formats the messages from the values by variable name.
func (p MessagesPlaceholder) FormatMessages(values map[string]any) ([]llms.ChatMessage, error) {
value, ok := values[p.VariableName]
if !ok {
return nil, fmt.Errorf("%w: %s should be a list of chat messages", ErrNeedChatMessageList, p.VariableName)
}
baseMessages, ok := value.([]llms.ChatMessage)
if !ok {
return nil, fmt.Errorf("%w: %s should be a list of chat messages", ErrNeedChatMessageList, p.VariableName)
}
return baseMessages, nil
}
// GetInputVariables returns the input variables the prompt expect.
func (p MessagesPlaceholder) GetInputVariables() []string {
return []string{p.VariableName}
}