llm/kimi/kimi.go

package kimi

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

type Kimi struct {
	userContext   *llm.UserContext
	browser       *rod.Browser
	launcher      *launcher.Launcher
	loginEvent    func(ctx context.Context, bytes []byte) error
	loginCallback func(ctx context.Context, user *llm.UserContext) error
}

func (k *Kimi) WithLoginEvent(f func(ctx context.Context, bytes []byte) error) llm.UserBrowserAgent {
	k.loginEvent = f
	return k
}

func (k *Kimi) WithLoginCallback(f func(ctx context.Context, user *llm.UserContext) error) llm.UserBrowserAgent {
	k.loginCallback = f
	return k
}

func (k *Kimi) WithUserContext(user *llm.UserContext) llm.UserBrowserAgent {
	k.userContext = user
	return k
}

func (k *Kimi) IsLogin(ctx context.Context, page *rod.Page) (bool, error) {
	_, err := page.Timeout(3 * time.Second).Element("div[data-testid='msh-header-user-avatar']")
	if err != nil && !errors.Is(err, context.DeadlineExceeded) {
		return false, err
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return false, nil
	}
	return true, nil
}

func (k *Kimi) Login(ctx context.Context, page *rod.Page) error {
	loginEle, err := page.Element("div[data-testid='msh-sidebar-user']")
	if err != nil {
		return err
	}
	err = loginEle.Click(proto.InputMouseButtonLeft, 1)
	if err != nil {
		return err
	}
	qrcodePreEle, err := page.ElementR("div p", "微信扫码登录")
	if err != nil {
		return err
	}
	qrcodeBox, err := qrcodePreEle.Parent()
	if err != nil {
		return err
	}
	screenshot, err := qrcodeBox.Screenshot(proto.PageCaptureScreenshotFormatPng, 10)
	if err != nil {
		return err
	}

	err = k.loginEvent(ctx, screenshot)
	if err != nil {
		return err
	}

	_, err = page.Element("div[data-testid='msh-header-user-avatar']")
	if err != nil {
		return err
	}

	localStorage := page.MustEval(`k => Object.assign({}, window.localStorage)`).Map()
	localStorageMap := make(map[string]string)
	for k, v := range localStorage {
		localStorageMap[k] = v.String()
	}
	err = k.loginCallback(ctx, llm.NewUserContext().WithLocalStorage(k.GetName(), localStorageMap).WithCookies(k.GetName(), proto.CookiesToParams(page.MustCookies())))
	if err != nil {
		return err
	}
	return nil
}

func (k *Kimi) GetName() string {
	return "kimi"
}

func (k *Kimi) NewChat(ctx context.Context) (llm.Chat, error) {
	c := &Chat{
		ai:       k,
		Page:     k.browser.Context(ctx).MustIncognito().MustPage(),
		qaSingle: make(chan bool),
	}
	wait := c.Page.MustWaitNavigation()

	// hook
	router := c.Page.HijackRequests()
	router.MustAdd("*.mp4", func(hijack *rod.Hijack) {
		if hijack.Request.Type() == proto.NetworkResourceTypeMedia {
			hijack.Response.Fail(proto.NetworkErrorReasonBlockedByClient)
			return
		}
		hijack.ContinueRequest(&proto.FetchContinueRequest{})
	})
	router.MustAdd("*/completion/stream", func(hijack *rod.Hijack) {
		hijack.MustLoadResponse()
		c.qaSingle <- true
	})
	go router.Run()

	err := c.Page.Navigate("https://kimi.moonshot.cn/")
	if err != nil {
		return nil, err
	}
	if k.userContext != nil {
		localStorage, ok := k.userContext.LocalStorage[k.GetName()]
		if ok {
			for key, val := range localStorage {
				c.Page.MustEval(fmt.Sprintf("() => {window.localStorage.setItem('%s','%s')}", key, val))
			}
		}
	}
	wait()
	return c, nil
}

func (k *Kimi) Close(ctx context.Context) error {
	k.launcher.Kill()
	return nil
}

func (k *Kimi) Init(ctx context.Context, debug bool) (llm.AI, error) {
	path, ok := launcher.LookPath()
	if !ok {
		return nil, errors.New("浏览器未找到")
	}
	launcherHandler := launcher.New().
		Bin(path).
		HeadlessNew(!debug).
		Set("disable-gpu").
		Devtools(debug)
	browser := rod.New().ControlURL(launcherHandler.MustLaunch()).MustConnect()
	return &Kimi{
		browser:  browser,
		launcher: launcherHandler,
	}, nil
}

func NewKimi() llm.AI {
	return &Kimi{}
}