Go使用Google Gemini Pro API创建简单聊天机器人

reader := bufio.NewReader(os.Stdin)
// 读取输入直到遇到第一个换行符
input, err := reader.ReadString('\n')

go get github.com/google/generative-ai-go

func getGeminiClient() *genai.Client {
  ctx := context.Background()
  client, err := genai.NewClient(ctx, option.WithAPIKey("YOUR-API-KEY"))
  if err != nil {
    log.Fatal(err)
  }
  return client
}

func handleChat(ctx context.Context, client *genai.Client, cs *genai.ChatSession, content string) error {
  if cs == nil {
    return errors.New("chat session is nil")
  }
  return geminiPro(ctx, cs, content)
}
func geminiPro(ctx context.Context, cs *genai.ChatSession, content string) error {
  prompt := genai.Text(content)
  iter := cs.SendMessageStream(ctx, prompt)
  fmt.Print("[Gemini]:")
  for {
    resp, err := iter.Next()
    if err == iterator.Done {
      break
    }
    if err != nil {
      fmt.Print("Something wrong, please try again later. \n")
      return err
    }
    printChatResponse(resp)
  }
  return nil
}
func printChatResponse(resp *genai.GenerateContentResponse) {
  for _, c := range resp.Candidates {
    for _, p := range c.Content.Parts {
      fmt.Printf("%v", p)
    }
  }
  fmt.Println()
}

// ResettableOnce 结构体包含一个互斥锁和一个标记是否已执行的布尔值
type ResettableOnce struct {
  m    sync.Mutex
  done bool
}
// Do 方法类似于 sync.Once 的 Do，但可以通过调用 Reset 来重置
func (o *ResettableOnce) Do(f func()) {
  o.m.Lock()
  defer o.m.Unlock()
  if !o.done {
    defer func() { o.done = true }()
    f()
  }
}
// Reset 方法重置 Once，允许再次调用 Do
func (o *ResettableOnce) Reset() {
  o.m.Lock()
  defer o.m.Unlock()
  o.done = false
}

func handleReset(client *genai.Client, cs **genai.ChatSession, history *[]*genai.Content) {
  fmt.Println("[System]", "Resetting chat session.")
  once.Reset()
  once.Do(func() {
    *cs = startChatSession(client, *history)
  })
}
func startChatSession(client *genai.Client, history []*genai.Content) *genai.ChatSession {
  model := client.GenerativeModel("gemini-pro")
  cs := model.StartChat()
  if len(history) > 0 {
    cs.History = history
  }
  return cs
}