-
Notifications
You must be signed in to change notification settings - Fork 1
/
messages.go
102 lines (91 loc) · 2.67 KB
/
messages.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
package main
import (
"fmt"
"github.com/charmbracelet/lipgloss"
"github.com/sashabaranov/go-openai"
)
const (
roleUser = "user"
roleSystem = "system"
modelUser = "userModel"
finishUser = "userEnd"
finishSystem = "systemEnd"
)
type (
Message struct {
Role string `json:"role"`
Content string `json:"content"`
FinishReason finishReason `json:"finish_reason"`
Model string `json:"name"` // WARN : for now it will mix the models and company
}
Conversation struct {
ID string `json:"id"`
Name string `json:"name"`
HasChange bool `json:"has_change"`
LastModel string `json:"last_model"`
Messages []Message `json:"messages"`
}
userOpenaiMessage openai.ChatCompletionMessage
gptMessage openai.ChatCompletionChoice
finishReason string
)
func (m Message) render() string {
senderStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("5"))
redStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("1"))
greenStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("2"))
blueStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("4"))
var style lipgloss.Style
var sender string
switch m.Role {
case roleUser:
sender = "You :"
case openai.ChatMessageRoleAssistant:
sender = "AI :"
case openai.ChatMessageRoleSystem:
sender = "System :"
}
switch m.FinishReason {
case finishUser:
style = senderStyle
case finishReason(openai.FinishReasonNull):
style = redStyle
case finishReason(openai.FinishReasonLength):
style = blueStyle
case finishReason(openai.FinishReasonStop):
style = greenStyle
}
return fmt.Sprintf("%s %s", style.Render(sender), m.Content)
}
func (conv *Conversation) openaiMessages() []openai.ChatCompletionMessage {
list := make([]openai.ChatCompletionMessage, len(conv.Messages))
for i, message := range conv.Messages {
list[i] = openai.ChatCompletionMessage{
Role: message.Role,
Content: message.Content,
}
}
return list
}
func (conv *Conversation) addMessage(message gptMessage, model string) {
// NOTE : we add a message only if there is a response
// We suppose that choices has a minimum size of 1
conv.Messages = append(conv.Messages, message.toMessage(model))
conv.HasChange = true
conv.LastModel = model
}
func (message userOpenaiMessage) toMessage() Message {
return Message{
Role: message.Role,
Content: message.Content,
FinishReason: "userEnd",
Model: "",
}
}
func (message gptMessage) toMessage(model string) Message {
return Message{
Role: message.Message.Role,
Content: message.Message.Content + "\n",
FinishReason: finishReason(message.FinishReason),
Model: model,
}
}