llm/kimi/chat.go

package kimi

import (
	"context"
	"errors"
	"git.ttmylife.com/zeuszhao/llm"
	"github.com/go-rod/rod"
	"github.com/go-rod/rod/lib/proto"
	"sync"
	"time"
)

type Chat struct {
	ai          llm.AI
	Page        *rod.Page
	qaSingle    chan bool
	userContext *llm.UserContext
	sync.Mutex
}

func (c *Chat) Ask(ctx context.Context, text string) (string, error) {
	c.TryLock()
	defer c.Unlock()
	ok, err := c.ai.IsLogin(ctx, c.Page)
	if err != nil {
		return "", err
	}
	if !ok {
		userContext, err := c.ai.Login(ctx, c.Page)
		if err != nil {
			return "", err
		}
		c.userContext = userContext
	}

	inputEle, err := c.Page.Element("#msh-chateditor")
	if err != nil {
		return "", err
	}

	err = inputEle.Input(text)
	if err != nil {
		return "", err
	}

	sendBtnEle, err := c.Page.Element("#send-button")
	if err != nil {
		return "", errors.Join(err, errors.New("send btn not fount"))
	}

	// 等待按钮可用
	err = sendBtnEle.WaitWritable()
	if err != nil {
		return "", errors.Join(err, errors.New("send btn not writable"))
	}

	err = sendBtnEle.WaitStable(1 * time.Second)
	if err != nil {
		return "", errors.Join(err, errors.New("send btn not stable"))
	}

	err = sendBtnEle.Click(proto.InputMouseButtonLeft, 1)
	if err != nil {
		return "", err
	}

	select {
	case <-c.qaSingle:
	case <-ctx.Done():
		return "", ctx.Err()
	case <-time.After(60 * time.Second):
		return "", errors.New("超时")
	}

	qaEleOk, err := c.Page.Element("button[data-testid='msh-chat-segment-reAnswer']")
	if err != nil {
		return "", err
	}
	err = qaEleOk.WaitVisible()
	if err != nil {
		return "", err
	}

	qaEle, err := c.Page.Elements(".pop-content")
	if err != nil {
		return "", err
	}
	return qaEle.Last().Text()
}

func (c *Chat) Close(ctx context.Context) error {
	return c.Page.Close()
}