-
Notifications
You must be signed in to change notification settings - Fork 0
/
common.go
105 lines (86 loc) · 2.19 KB
/
common.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
// Package gemini provides communication with Gemini.
package gemini
import (
"context"
"fmt"
"os"
"path/filepath"
"strings"
"github.com/google/generative-ai-go/genai"
"google.golang.org/api/iterator"
"google.golang.org/api/option"
)
type answerForStream struct {
Answer string
Err error
}
type chatHistory struct {
Role string
Answer string
}
func responseToAnswer(response *genai.GenerateContentResponse) string {
if response == nil {
return ""
}
answer := ""
for _, candidate := range response.Candidates {
answer += contentToAnswer(candidate.Content)
}
return answer
}
func contentToAnswer(content *genai.Content) string {
if content == nil {
return ""
}
answer := ""
for _, part := range content.Parts {
builder := strings.Builder{}
fmt.Fprintf(&builder, "%v", part)
answer += builder.String()
}
return answer
}
func makeParts(text string, images []string) ([]genai.Part, error) {
parts := []genai.Part{}
if len(text) != 0 {
parts = append(parts, genai.Text(text))
}
for _, image := range images {
if data, err := os.ReadFile(image); err != nil {
return nil, err
} else if extension := strings.ToLower(filepath.Ext(image)); len(extension) > 0 && extension[0] == '.' {
parts = append(parts, genai.ImageData(extension[1:], data))
} else {
parts = append(parts, genai.ImageData(extension, data))
}
}
return parts, nil
}
func getClient(ctx context.Context, apiKey string) (*genai.Client, error) {
return genai.NewClient(ctx, option.WithAPIKey(apiKey))
}
func getModel(client *genai.Client, images []string) *genai.GenerativeModel {
if len(images) == 0 {
return client.GenerativeModel("gemini-pro")
} else {
return client.GenerativeModel("gemini-pro-vision")
}
}
func getStreamChannel(responseIterator *genai.GenerateContentResponseIterator) chan answerForStream {
channel := make(chan answerForStream)
go func() {
for {
if response, err := responseIterator.Next(); err == iterator.Done {
close(channel)
break
} else if err != nil {
channel <- answerForStream{Answer: "", Err: err}
close(channel)
break
} else {
channel <- answerForStream{Answer: responseToAnswer(response), Err: nil}
}
}
}()
return channel
}