137 lines
3.3 KiB
Go
137 lines
3.3 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
|
|
"go.wit.com/lib/protobuf/chatpb"
|
|
"go.wit.com/log"
|
|
"google.golang.org/genai"
|
|
)
|
|
|
|
func initGeminiAPI() error {
|
|
if me.ctx != nil {
|
|
// already initialized
|
|
return nil
|
|
}
|
|
apiKey := os.Getenv("GEMINI_API_KEY")
|
|
if apiKey == "" {
|
|
return log.Errorf("GEMINI_API_KEY environment variable not set")
|
|
}
|
|
|
|
me.ctx = context.Background()
|
|
var err error
|
|
me.client, err = genai.NewClient(me.ctx, &genai.ClientConfig{APIKey: apiKey})
|
|
if err != nil {
|
|
return log.Errorf("failed to create new genai client: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// doConnect initializes the Gemini client and handles the request flow.
|
|
func doConnect() error {
|
|
initGeminiAPI()
|
|
|
|
if me.lastChat == nil {
|
|
log.Info("WTF. lastChat is nil")
|
|
return nil
|
|
}
|
|
|
|
// if me.lastChat.Entries == nil {
|
|
// me.lastChat.Entries = new(chatpb.ChatEntry)
|
|
// }
|
|
|
|
// In a real application, you would get user input here.
|
|
// For now, we'll use a hardcoded prompt.
|
|
if len(me.lastChat.GetEntries()) == 0 {
|
|
me.lastChat.Entries = append(me.lastChat.Entries, &chatpb.ChatEntry{
|
|
Parts: []*chatpb.Part{
|
|
{PartType: &chatpb.Part_Text{Text: "hello, how are you"}},
|
|
},
|
|
})
|
|
}
|
|
|
|
lastEntry := me.lastChat.GetEntries()[len(me.lastChat.GetEntries())-1]
|
|
genaiContents, err := convertToGenai(lastEntry.GetGeminiRequest())
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
resp, err := me.client.Models.GenerateContent(me.ctx, "gemini-2.5-flash", genaiContents, nil)
|
|
if err != nil {
|
|
return log.Errorf("error sending message: %v", err)
|
|
}
|
|
|
|
if resp == nil || len(resp.Candidates) == 0 || resp.Candidates[0].Content == nil {
|
|
log.Info("Received an empty response from the API. Stopping.")
|
|
return nil
|
|
}
|
|
|
|
// Append the model's response to the history
|
|
me.lastChat.Entries = append(me.lastChat.Entries, convertToPB(resp))
|
|
|
|
// Check for a function call
|
|
hasFunctionCall := false
|
|
for _, part := range resp.Candidates[0].Content.Parts {
|
|
if fc := part.FunctionCall; fc != nil {
|
|
hasFunctionCall = true
|
|
functionResponse := handleFunctionCall(fc)
|
|
// Append the function response to the history for the next turn
|
|
me.lastChat.Entries = append(me.lastChat.Entries, &chatpb.ChatEntry{
|
|
Parts: []*chatpb.Part{
|
|
{PartType: &chatpb.Part_FunctionResponse{
|
|
FunctionResponse: &chatpb.FunctionResponse{
|
|
Name: functionResponse.Name,
|
|
// TODO: map response
|
|
},
|
|
}},
|
|
},
|
|
})
|
|
}
|
|
}
|
|
|
|
// If there was no function call, print the text and stop.
|
|
if !hasFunctionCall {
|
|
log.Info("Response from API:")
|
|
for _, cand := range resp.Candidates {
|
|
if cand.Content != nil {
|
|
for _, part := range cand.Content.Parts {
|
|
if part.Text != "" {
|
|
fmt.Println(part.Text)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// sampleHello sends a hardcoded prompt to the model and prints the response.
|
|
func simpleHello() error {
|
|
log.Info("Sending 'hello, how are you' to the Gemini API...")
|
|
|
|
// Create the parts slice
|
|
parts := []*genai.Part{
|
|
{Text: "What is my brothers name?"},
|
|
}
|
|
|
|
content := []*genai.Content{{Parts: parts}}
|
|
|
|
resp, err := me.client.Models.GenerateContent(me.ctx, "gemini-2.5-flash", content, nil)
|
|
if err != nil {
|
|
return log.Errorf("error sending message: %v", err)
|
|
}
|
|
|
|
log.Info("Response from API:")
|
|
for _, cand := range resp.Candidates {
|
|
if cand.Content != nil {
|
|
for _, part := range cand.Content.Parts {
|
|
fmt.Println(part)
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|