mirror of
https://github.com/googleapis/genai-toolbox.git
synced 2026-02-11 23:55:07 -05:00
175 lines
5.3 KiB
Go
175 lines
5.3 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
|
|
"github.com/googleapis/mcp-toolbox-sdk-go/core"
|
|
"google.golang.org/genai"
|
|
)
|
|
|
|
// ConvertToGenaiTool translates a ToolboxTool into the genai.FunctionDeclaration format.
|
|
func ConvertToGenaiTool(toolboxTool *core.ToolboxTool) *genai.Tool {
|
|
|
|
inputschema, err := toolboxTool.InputSchema()
|
|
if err != nil {
|
|
return &genai.Tool{}
|
|
}
|
|
|
|
var paramsSchema *genai.Schema
|
|
_ = json.Unmarshal(inputschema, ¶msSchema)
|
|
// First, create the function declaration.
|
|
funcDeclaration := &genai.FunctionDeclaration{
|
|
Name: toolboxTool.Name(),
|
|
Description: toolboxTool.Description(),
|
|
Parameters: paramsSchema,
|
|
}
|
|
|
|
// Then, wrap the function declaration in a genai.Tool struct.
|
|
return &genai.Tool{
|
|
FunctionDeclarations: []*genai.FunctionDeclaration{funcDeclaration},
|
|
}
|
|
}
|
|
|
|
func printResponse(resp *genai.GenerateContentResponse) {
|
|
for _, cand := range resp.Candidates {
|
|
if cand.Content != nil {
|
|
for _, part := range cand.Content.Parts {
|
|
fmt.Println(part.Text)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
const systemPrompt = `
|
|
You're a helpful hotel assistant. You handle hotel searching, booking, and
|
|
cancellations. When the user searches for a hotel, mention its name, id,
|
|
location and price tier. Always mention hotel ids while performing any
|
|
searches. This is very important for any operations. For any bookings or
|
|
cancellations, please provide the appropriate confirmation. Be sure to
|
|
update checkin or checkout dates if mentioned by the user.
|
|
Don't ask for confirmations from the user.
|
|
`
|
|
|
|
var queries = []string{
|
|
"Find hotels in Basel with Basel in its name.",
|
|
"Can you book the hotel Hilton Basel for me?",
|
|
"Oh wait, this is too expensive. Please cancel it.",
|
|
"Please book the Hyatt Regency instead.",
|
|
"My check in dates would be from April 10, 2024 to April 19, 2024.",
|
|
}
|
|
|
|
func main() {
|
|
// Setup
|
|
ctx := context.Background()
|
|
apiKey := os.Getenv("GOOGLE_API_KEY")
|
|
toolboxURL := "http://localhost:5000"
|
|
|
|
// Initialize the Google GenAI client using the explicit ClientConfig.
|
|
client, err := genai.NewClient(ctx, &genai.ClientConfig{
|
|
APIKey: apiKey,
|
|
})
|
|
if err != nil {
|
|
log.Fatalf("Failed to create Google GenAI client: %v", err)
|
|
}
|
|
|
|
// Initialize the MCP Toolbox client.
|
|
toolboxClient, err := core.NewToolboxClient(toolboxURL)
|
|
if err != nil {
|
|
log.Fatalf("Failed to create Toolbox client: %v", err)
|
|
}
|
|
|
|
// Load the tool using the MCP Toolbox SDK.
|
|
tools, err := toolboxClient.LoadToolset("my-toolset", ctx)
|
|
if err != nil {
|
|
log.Fatalf("Failed to load tools: %v\nMake sure your Toolbox server is running and the tool is configured.", err)
|
|
}
|
|
|
|
genAITools := make([]*genai.Tool, len(tools))
|
|
toolsMap := make(map[string]*core.ToolboxTool, len(tools))
|
|
|
|
for i, tool := range tools {
|
|
genAITools[i] = ConvertToGenaiTool(tool)
|
|
toolsMap[tool.Name()] = tool
|
|
}
|
|
|
|
// Set up the generative model with the available tool.
|
|
modelName := "gemini-2.0-flash"
|
|
|
|
// Create the initial content prompt for the model.
|
|
messageHistory := []*genai.Content{
|
|
genai.NewContentFromText(systemPrompt, genai.RoleUser),
|
|
}
|
|
config := &genai.GenerateContentConfig{
|
|
Tools: genAITools,
|
|
ToolConfig: &genai.ToolConfig{
|
|
FunctionCallingConfig: &genai.FunctionCallingConfig{
|
|
Mode: genai.FunctionCallingConfigModeAny,
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, query := range queries {
|
|
|
|
messageHistory = append(messageHistory, genai.NewContentFromText(query, genai.RoleUser))
|
|
|
|
genContentResp, err := client.Models.GenerateContent(ctx, modelName, messageHistory, config)
|
|
if err != nil {
|
|
log.Fatalf("LLM call failed for query '%s': %v", query, err)
|
|
}
|
|
|
|
if len(genContentResp.Candidates) > 0 && genContentResp.Candidates[0].Content != nil {
|
|
messageHistory = append(messageHistory, genContentResp.Candidates[0].Content)
|
|
}
|
|
|
|
functionCalls := genContentResp.FunctionCalls()
|
|
|
|
toolResponseParts := []*genai.Part{}
|
|
|
|
for _, fc := range functionCalls {
|
|
|
|
toolToInvoke, found := toolsMap[fc.Name]
|
|
if !found {
|
|
log.Fatalf("Tool '%s' not found in loaded tools map. Check toolset configuration.", fc.Name)
|
|
}
|
|
|
|
toolResult, invokeErr := toolToInvoke.Invoke(ctx, fc.Args)
|
|
if invokeErr != nil {
|
|
log.Fatalf("Failed to execute tool '%s': %v", fc.Name, invokeErr)
|
|
}
|
|
|
|
// Enhanced Tool Result Handling (retained to prevent nil issues)
|
|
toolResultString := ""
|
|
if toolResult != nil {
|
|
jsonBytes, marshalErr := json.Marshal(toolResult)
|
|
if marshalErr == nil {
|
|
toolResultString = string(jsonBytes)
|
|
} else {
|
|
toolResultString = fmt.Sprintf("%v", toolResult)
|
|
}
|
|
}
|
|
|
|
responseMap := map[string]any{"result": toolResultString}
|
|
|
|
toolResponseParts = append(toolResponseParts, genai.NewPartFromFunctionResponse(fc.Name, responseMap))
|
|
}
|
|
// Add all accumulated tool responses for this turn to the message history.
|
|
toolResponseContent := genai.NewContentFromParts(toolResponseParts, "function")
|
|
messageHistory = append(messageHistory, toolResponseContent)
|
|
|
|
finalResponse, err := client.Models.GenerateContent(ctx, modelName, messageHistory, &genai.GenerateContentConfig{})
|
|
if err != nil {
|
|
log.Fatalf("Error calling GenerateContent (with function result): %v", err)
|
|
}
|
|
|
|
printResponse(finalResponse)
|
|
// Add the final textual response from the LLM to the history
|
|
if len(finalResponse.Candidates) > 0 && finalResponse.Candidates[0].Content != nil {
|
|
messageHistory = append(messageHistory, finalResponse.Candidates[0].Content)
|
|
}
|
|
}
|
|
}
|