Skip to content

Commit 5c8b73c

Browse files
elibeneliben
authored
vertexai: add NewUserContent helper function (#10570)
1 parent ba82942 commit 5c8b73c

File tree

6 files changed

+21
-19
lines changed

6 files changed

+21
-19
lines changed

‎vertexai/genai/caching_test.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -103,9 +103,9 @@ func testCaching(t *testing.T, client *Client) {
103103
argcc := &CachedContent{
104104
Model: model,
105105
Expiration: ExpireTimeOrTTL{TTL: ttl},
106-
Contents: []*Content{{Role: "user", Parts: []Part{
107-
FileData{MIMEType: "text/plain", FileURI: gcsFilePath},
108-
}}},
106+
Contents: []*Content{NewUserContent(FileData{
107+
MIMEType: "text/plain",
108+
FileURI: gcsFilePath})},
109109
}
110110
cc := must(client.CreateCachedContent(ctx, argcc))
111111
compare(cc, wantExpireTime)

‎vertexai/genai/chat.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ func (m *GenerativeModel) StartChat() *ChatSession {
3232
// SendMessage sends a request to the model as part of a chat session.
3333
func (cs *ChatSession) SendMessage(ctx context.Context, parts ...Part) (*GenerateContentResponse, error) {
3434
// Call the underlying client with the entire history plus the argument Content.
35-
cs.History = append(cs.History, newUserContent(parts))
35+
cs.History = append(cs.History, NewUserContent(parts...))
3636
req := cs.m.newGenerateContentRequest(cs.History...)
3737
cc := int32(1)
3838
req.GenerationConfig.CandidateCount = &cc
@@ -46,7 +46,7 @@ func (cs *ChatSession) SendMessage(ctx context.Context, parts ...Part) (*Generat
4646

4747
// SendMessageStream is like SendMessage, but with a streaming request.
4848
func (cs *ChatSession) SendMessageStream(ctx context.Context, parts ...Part) *GenerateContentResponseIterator {
49-
cs.History = append(cs.History, newUserContent(parts))
49+
cs.History = append(cs.History, NewUserContent(parts...))
5050
req := cs.m.newGenerateContentRequest(cs.History...)
5151
var cc int32 = 1
5252
req.GenerationConfig.CandidateCount = &cc

‎vertexai/genai/client.go

+3-7
Original file line numberDiff line numberDiff line change
@@ -187,12 +187,12 @@ func (m *GenerativeModel) Name() string {
187187

188188
// GenerateContent produces a single request and response.
189189
func (m *GenerativeModel) GenerateContent(ctx context.Context, parts ...Part) (*GenerateContentResponse, error) {
190-
return m.generateContent(ctx, m.newGenerateContentRequest(newUserContent(parts)))
190+
return m.generateContent(ctx, m.newGenerateContentRequest(NewUserContent(parts...)))
191191
}
192192

193193
// GenerateContentStream returns an iterator that enumerates responses.
194194
func (m *GenerativeModel) GenerateContentStream(ctx context.Context, parts ...Part) *GenerateContentResponseIterator {
195-
streamClient, err := m.c.pc.StreamGenerateContent(ctx, m.newGenerateContentRequest(newUserContent(parts)))
195+
streamClient, err := m.c.pc.StreamGenerateContent(ctx, m.newGenerateContentRequest(NewUserContent(parts...)))
196196
return &GenerateContentResponseIterator{
197197
sc: streamClient,
198198
err: err,
@@ -221,10 +221,6 @@ func (m *GenerativeModel) newGenerateContentRequest(contents ...*Content) *pb.Ge
221221
}
222222
}
223223

224-
func newUserContent(parts []Part) *Content {
225-
return &Content{Role: roleUser, Parts: parts}
226-
}
227-
228224
// GenerateContentResponseIterator is an iterator over GnerateContentResponse.
229225
type GenerateContentResponseIterator struct {
230226
sc pb.PredictionService_StreamGenerateContentClient
@@ -286,7 +282,7 @@ func protoToResponse(resp *pb.GenerateContentResponse) (*GenerateContentResponse
286282

287283
// CountTokens counts the number of tokens in the content.
288284
func (m *GenerativeModel) CountTokens(ctx context.Context, parts ...Part) (*CountTokensResponse, error) {
289-
req := m.newCountTokensRequest(newUserContent(parts))
285+
req := m.newCountTokensRequest(NewUserContent(parts...))
290286
res, err := m.c.pc.CountTokens(ctx, req)
291287
if err != nil {
292288
return nil, err

‎vertexai/genai/client_test.go

+1-3
Original file line numberDiff line numberDiff line change
@@ -63,9 +63,7 @@ func TestLive(t *testing.T) {
6363
t.Run("system-instructions", func(t *testing.T) {
6464
model := client.GenerativeModel(defaultModel)
6565
model.Temperature = Ptr[float32](0)
66-
model.SystemInstruction = &Content{
67-
Parts: []Part{Text("You are Yoda from Star Wars.")},
68-
}
66+
model.SystemInstruction = NewUserContent(Text("You are Yoda from Star Wars."))
6967
resp, err := model.GenerateContent(ctx, Text("What is the average size of a swallow?"))
7068
if err != nil {
7169
t.Fatal(err)

‎vertexai/genai/content.go

+10
Original file line numberDiff line numberDiff line change
@@ -148,3 +148,13 @@ func (c *Candidate) FunctionCalls() []FunctionCall {
148148
}
149149
return fcs
150150
}
151+
152+
// NewUserContent returns a [Content] with a "user" role set and one or more
153+
// parts.
154+
func NewUserContent(parts ...Part) *Content {
155+
content := &Content{Role: roleUser, Parts: []Part{}}
156+
for _, part := range parts {
157+
content.Parts = append(content.Parts, part)
158+
}
159+
return content
160+
}

‎vertexai/genai/example_test.go

+2-4
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,7 @@ func ExampleGenerativeModel_GenerateContent_config() {
7474
model.SetTopP(0.5)
7575
model.SetTopK(20)
7676
model.SetMaxOutputTokens(100)
77-
model.SystemInstruction = &genai.Content{
78-
Parts: []genai.Part{genai.Text("You are Yoda from Star Wars.")},
79-
}
77+
model.SystemInstruction = genai.NewUserContent(genai.Text("You are Yoda from Star Wars."))
8078
resp, err := model.GenerateContent(ctx, genai.Text("What is the average size of a swallow?"))
8179
if err != nil {
8280
log.Fatal(err)
@@ -315,7 +313,7 @@ func ExampleClient_cachedContent() {
315313
file := genai.FileData{MIMEType: "application/pdf", FileURI: "gs://my-bucket/my-doc.pdf"}
316314
cc, err := client.CreateCachedContent(ctx, &genai.CachedContent{
317315
Model: modelName,
318-
Contents: []*genai.Content{{Parts: []genai.Part{file}}},
316+
Contents: []*genai.Content{genai.NewUserContent(file)},
319317
})
320318
model := client.GenerativeModelFromCachedContent(cc)
321319
// Work with the model as usual in this program.

0 commit comments

Comments
 (0)