1

.idea/.gitignore

@@ -0,0 +1,8 @@
+# 默认忽略的文件
+/shelf/
+/workspace.xml
+# 基于编辑器的 HTTP 客户端请求
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

go.mod

@@ -0,0 +1,7 @@
+module git.ttmylife.com/zeuszhao/llm
+
+go 1.23.1
+
+require github.com/go-rod/rod v0.116.2
+
+require github.com/ysmood/gson v0.7.3 // indirect

go.sum

@@ -0,0 +1,8 @@
+github.com/go-rod/rod v0.116.2 h1:A5t2Ky2A+5eD/ZJQr1EfsQSe5rms5Xof/qj296e+ZqA=
+github.com/go-rod/rod v0.116.2/go.mod h1:H+CMO9SCNc2TJ2WfrG+pKhITz57uGNYU43qYHh438Mg=
+github.com/ysmood/gop v0.2.0 h1:+tFrG0TWPxT6p9ZaZs+VY+opCvHU8/3Fk6BaNv6kqKg=
+github.com/ysmood/gop v0.2.0/go.mod h1:rr5z2z27oGEbyB787hpEcx4ab8cCiPnKxn0SUHt6xzk=
+github.com/ysmood/got v0.40.0 h1:ZQk1B55zIvS7zflRrkGfPDrPG3d7+JOza1ZkNxcc74Q=
+github.com/ysmood/got v0.40.0/go.mod h1:W7DdpuX6skL3NszLmAsC5hT7JAhuLZhByVzHTq874Qg=
+github.com/ysmood/gson v0.7.3 h1:QFkWbTH8MxyUTKPkVWAENJhxqdBa4lYTQWqZCiLG6kE=
+github.com/ysmood/gson v0.7.3/go.mod h1:3Kzs5zDl21g5F/BlLTNcuAGAYLKt2lV5G8D1zF3RNmg=

kimi/chat.go

@@ -0,0 +1,91 @@
+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
+	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 {
+		err = c.ai.Login(ctx, c.Page)
+		if err != nil {
+			return "", err
+		}
+	}
+
+	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()
+}

kimi/kimi.go

@@ -0,0 +1,151 @@
+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 {
+	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) 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, user *llm.UserContext) (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
+	}
+	localStorage, ok := user.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) (llm.AI, error) {
+	path, ok := launcher.LookPath()
+	if !ok {
+		return nil, errors.New("浏览器未找到")
+	}
+	launcherHandler := launcher.New().
+		Bin(path).
+		HeadlessNew(true).
+		Set("disable-gpu").
+		Devtools(false)
+	browser := rod.New().ControlURL(launcherHandler.MustLaunch()).MustConnect()
+	return &Kimi{
+		browser:  browser,
+		launcher: launcherHandler,
+	}, nil
+}
+
+func NewKimi() llm.AI {
+	return &Kimi{}
+}

llm.go

@@ -0,0 +1,17 @@
+package llm
+
+import (
+	"context"
+)
+
+type Chat interface {
+	Ask(ctx context.Context, text string) (string, error)
+	Close(ctx context.Context) error
+}
+
+type AI interface {
+	UserBrowserAgent
+	Init(ctx context.Context) (AI, error)
+	NewChat(ctx context.Context, user *UserContext) (Chat, error)
+	Close(ctx context.Context) error
+}

user.go

@@ -0,0 +1,58 @@
+package llm
+
+import (
+	"context"
+	"encoding/json"
+	"github.com/go-rod/rod"
+	"github.com/go-rod/rod/lib/proto"
+	"sync"
+)
+
+type UserBrowserAgent interface {
+	IsLogin(ctx context.Context, page *rod.Page) (bool, error)
+	Login(ctx context.Context, page *rod.Page) error
+	WithLoginEvent(f func(ctx context.Context, bytes []byte) error) UserBrowserAgent
+	WithLoginCallback(f func(ctx context.Context, user *UserContext) error) UserBrowserAgent
+	GetName() string
+}
+
+type UserContext struct {
+	Cookies      map[string][]*proto.NetworkCookieParam `json:"cookies"`
+	LocalStorage map[string]map[string]string           `json:"local_storage"`
+	sync.Mutex
+}
+
+func NewUserContext() *UserContext {
+	return &UserContext{
+		Cookies:      make(map[string][]*proto.NetworkCookieParam),
+		LocalStorage: make(map[string]map[string]string),
+	}
+}
+
+func (uc *UserContext) WithCookies(name string, cookies []*proto.NetworkCookieParam) *UserContext {
+	uc.Cookies[name] = cookies
+	return uc
+}
+
+func (uc *UserContext) WithLocalStorage(name string, l map[string]string) *UserContext {
+	uc.LocalStorage[name] = l
+	return uc
+}
+
+// ToString 将用户上下文转换为字符串
+func (uc *UserContext) ToString() (string, error) {
+	byt, err := json.Marshal(uc)
+	return string(byt), err
+}
+
+// ParseContextFromJson 解析json字符串,设置Cookie
+func (uc *UserContext) ParseContextFromJson(jsonStr string) error {
+	uc1 := &UserContext{}
+	err := json.Unmarshal([]byte(jsonStr), uc1)
+	if err != nil {
+		return err
+	}
+	uc.Cookies = uc1.Cookies
+	uc.LocalStorage = uc1.LocalStorage
+	return nil
+}