import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"cloud.google.com/go/vertexai/genai"
)
// functionCallsChat opens a chat session and sends 4 messages to the model:
// - convert a first text question into a structured function call request
// - convert the first structured function call response into natural language
// - convert a second text question into a structured function call request
// - convert the second structured function call response into natural language
func functionCallsChat(w io.Writer, projectID, location, modelName string) error {
// location := "us-central1"
// modelName := "gemini-1.5-flash-001"
ctx := context.Background()
client, err := genai.NewClient(ctx, projectID, location)
if err != nil {
return fmt.Errorf("unable to create client: %w", err)
}
defer client.Close()
model := client.GenerativeModel(modelName)
// Build an OpenAPI schema, in memory
paramsProduct := &genai.Schema{
Type: genai.TypeObject,
Properties: map[string]*genai.Schema{
"productName": {
Type: genai.TypeString,
Description: "Product name",
},
},
}
fundeclProductInfo := &genai.FunctionDeclaration{
Name: "getProductSku",
Description: "Get the SKU for a product",
Parameters: paramsProduct,
}
paramsStore := &genai.Schema{
Type: genai.TypeObject,
Properties: map[string]*genai.Schema{
"location": {
Type: genai.TypeString,
Description: "Location",
},
},
}
fundeclStoreLocation := &genai.FunctionDeclaration{
Name: "getStoreLocation",
Description: "Get the location of the closest store",
Parameters: paramsStore,
}
model.Tools = []*genai.Tool{
{FunctionDeclarations: []*genai.FunctionDeclaration{
fundeclProductInfo,
fundeclStoreLocation,
}},
}
model.SetTemperature(0.0)
chat := model.StartChat()
// Send a prompt for the first conversation turn that should invoke the getProductSku function
prompt := "Do you have the Pixel 8 Pro in stock?"
fmt.Fprintf(w, "Question: %s\n", prompt)
resp, err := chat.SendMessage(ctx, genai.Text(prompt))
if err != nil {
return err
}
if len(resp.Candidates) == 0 ||
len(resp.Candidates[0].Content.Parts) == 0 {
return errors.New("empty response from model")
}
// The model has returned a function call to the declared function `getProductSku`
// with a value for the argument `productName`.
jsondata, err := json.MarshalIndent(resp.Candidates[0].Content.Parts[0], "\t", " ")
if err != nil {
return fmt.Errorf("json.MarshalIndent: %w", err)
}
fmt.Fprintf(w, "function call generated by the model:\n\t%s\n", string(jsondata))
// Create a function call response, to simulate the result of a call to a
// real service
funresp := &genai.FunctionResponse{
Name: "getProductSku",
Response: map[string]any{
"sku": "GA04834-US",
"in_stock": "yes",
},
}
jsondata, err = json.MarshalIndent(funresp, "\t", " ")
if err != nil {
return fmt.Errorf("json.MarshalIndent: %w", err)
}
fmt.Fprintf(w, "function call response sent to the model:\n\t%s\n\n", string(jsondata))
// And provide the function call response to the model
resp, err = chat.SendMessage(ctx, funresp)
if err != nil {
return err
}
if len(resp.Candidates) == 0 ||
len(resp.Candidates[0].Content.Parts) == 0 {
return errors.New("empty response from model")
}
// The model has taken the function call response as input, and has
// reformulated the response to the user.
jsondata, err = json.MarshalIndent(resp.Candidates[0].Content.Parts[0], "\t", " ")
if err != nil {
return fmt.Errorf("json.MarshalIndent: %w", err)
}
fmt.Fprintf(w, "Answer generated by the model:\n\t%s\n\n", string(jsondata))
// Send a prompt for the second conversation turn that should invoke the getStoreLocation function
prompt2 := "Is there a store in Mountain View, CA that I can visit to try it out?"
fmt.Fprintf(w, "Question: %s\n", prompt)
resp, err = chat.SendMessage(ctx, genai.Text(prompt2))
if err != nil {
return err
}
if len(resp.Candidates) == 0 ||
len(resp.Candidates[0].Content.Parts) == 0 {
return errors.New("empty response from model")
}
// The model has returned a function call to the declared function `getStoreLocation`
// with a value for the argument `store`.
jsondata, err = json.MarshalIndent(resp.Candidates[0].Content.Parts[0], "\t", " ")
if err != nil {
return fmt.Errorf("json.MarshalIndent: %w", err)
}
fmt.Fprintf(w, "function call generated by the model:\n\t%s\n", string(jsondata))
// Create a function call response, to simulate the result of a call to a
// real service
funresp = &genai.FunctionResponse{
Name: "getStoreLocation",
Response: map[string]any{
"store": "2000 N Shoreline Blvd, Mountain View, CA 94043, US",
},
}
jsondata, err = json.MarshalIndent(funresp, "\t", " ")
if err != nil {
return fmt.Errorf("json.MarshalIndent: %w", err)
}
fmt.Fprintf(w, "function call response sent to the model:\n\t%s\n\n", string(jsondata))
// And provide the function call response to the model
resp, err = chat.SendMessage(ctx, funresp)
if err != nil {
return err
}
if len(resp.Candidates) == 0 ||
len(resp.Candidates[0].Content.Parts) == 0 {
return errors.New("empty response from model")
}
// The model has taken the function call response as input, and has
// reformulated the response to the user.
jsondata, err = json.MarshalIndent(resp.Candidates[0].Content.Parts[0], "\t", " ")
if err != nil {
return fmt.Errorf("json.MarshalIndent: %w", err)
}
fmt.Fprintf(w, "Answer generated by the model:\n\t%s\n\n", string(jsondata))
return nil
}
