llm/user.go

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) (*UserContext, error)
	WithLoginEvent(f func(ctx context.Context, bytes []byte) error) UserBrowserAgent
	WithLoginCallback(f func(ctx context.Context, user *UserContext) error) UserBrowserAgent
}

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
}