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
|
package main
import (
"bytes"
"elefant/models"
"encoding/json"
"io"
"strings"
)
type ChunkParser interface {
ParseChunk([]byte) (string, bool, error)
FormMsg(msg, role string) (io.Reader, error)
}
func initChunkParser() {
chunkParser = LlamaCPPeer{}
if strings.Contains(cfg.CurrentAPI, "v1") {
logger.Info("chosen openai parser")
chunkParser = OpenAIer{}
return
}
logger.Info("chosen llamacpp parser")
}
type LlamaCPPeer struct {
}
type OpenAIer struct {
}
func (lcp LlamaCPPeer) FormMsg(msg, role string) (io.Reader, error) {
if msg != "" { // otherwise let the bot continue
// if role == cfg.UserRole {
// msg = msg + cfg.AssistantRole + ":"
// }
newMsg := models.RoleMsg{Role: role, Content: msg}
chatBody.Messages = append(chatBody.Messages, newMsg)
// if rag
if cfg.RAGEnabled {
ragResp, err := chatRagUse(newMsg.Content)
if err != nil {
logger.Error("failed to form a rag msg", "error", err)
return nil, err
}
ragMsg := models.RoleMsg{Role: cfg.ToolRole, Content: ragResp}
chatBody.Messages = append(chatBody.Messages, ragMsg)
}
}
messages := make([]string, len(chatBody.Messages))
for i, m := range chatBody.Messages {
messages[i] = m.ToPrompt()
}
prompt := strings.Join(messages, "\n")
botMsgStart := "\n" + cfg.AssistantRole + ":\n"
payload := models.NewLCPReq(prompt+botMsgStart, role, defaultLCPProps)
data, err := json.Marshal(payload)
if err != nil {
logger.Error("failed to form a msg", "error", err)
return nil, err
}
return bytes.NewReader(data), nil
}
func (lcp LlamaCPPeer) ParseChunk(data []byte) (string, bool, error) {
llmchunk := models.LlamaCPPResp{}
if err := json.Unmarshal(data, &llmchunk); err != nil {
logger.Error("failed to decode", "error", err, "line", string(data))
return "", false, err
}
if llmchunk.Stop {
if llmchunk.Content != "" {
logger.Error("text inside of finish llmchunk", "chunk", llmchunk)
}
return llmchunk.Content, true, nil
}
return llmchunk.Content, false, nil
}
func (op OpenAIer) ParseChunk(data []byte) (string, bool, error) {
llmchunk := models.LLMRespChunk{}
if err := json.Unmarshal(data, &llmchunk); err != nil {
logger.Error("failed to decode", "error", err, "line", string(data))
return "", false, err
}
content := llmchunk.Choices[len(llmchunk.Choices)-1].Delta.Content
if llmchunk.Choices[len(llmchunk.Choices)-1].FinishReason == "stop" {
if content != "" {
logger.Error("text inside of finish llmchunk", "chunk", llmchunk)
}
return content, true, nil
}
return content, false, nil
}
func (op OpenAIer) FormMsg(msg, role string) (io.Reader, error) {
if msg != "" { // otherwise let the bot continue
newMsg := models.RoleMsg{Role: role, Content: msg}
chatBody.Messages = append(chatBody.Messages, newMsg)
// if rag
if cfg.RAGEnabled {
ragResp, err := chatRagUse(newMsg.Content)
if err != nil {
logger.Error("failed to form a rag msg", "error", err)
return nil, err
}
ragMsg := models.RoleMsg{Role: cfg.ToolRole, Content: ragResp}
chatBody.Messages = append(chatBody.Messages, ragMsg)
}
}
data, err := json.Marshal(chatBody)
if err != nil {
logger.Error("failed to form a msg", "error", err)
return nil, err
}
return bytes.NewReader(data), nil
}
|
