4

kimi/chat.go

@@ -11,9 +11,10 @@ import (
 )
 
 type Chat struct {
-	ai       llm.AI
-	Page     *rod.Page
-	qaSingle chan bool
+	ai          llm.AI
+	Page        *rod.Page
+	qaSingle    chan bool
+	userContext *llm.UserContext
 	sync.Mutex
 }
 
@@ -25,10 +26,11 @@ func (c *Chat) Ask(ctx context.Context, text string) (string, error) {
 		return "", err
 	}
 	if !ok {
-		err = c.ai.Login(ctx, c.Page)
+		userContext, err := c.ai.Login(ctx, c.Page)
 		if err != nil {
 			return "", err
 		}
+		c.userContext = userContext
 	}
 
 	inputEle, err := c.Page.Element("#msh-chateditor")
@@ -89,3 +91,8 @@ func (c *Chat) Ask(ctx context.Context, text string) (string, error) {
 func (c *Chat) Close(ctx context.Context) error {
 	return c.Page.Close()
 }
+
+func (c *Chat) WithUserContext(user *llm.UserContext) llm.Chat {
+	c.userContext = user
+	return c
+}

kimi/kimi.go

@@ -12,7 +12,6 @@ import (
 )
 
 type Kimi struct {
-	userContext   *llm.UserContext
 	browser       *rod.Browser
 	launcher      *launcher.Launcher
 	loginEvent    func(ctx context.Context, bytes []byte) error
@@ -29,11 +28,6 @@ func (k *Kimi) WithLoginCallback(f func(ctx context.Context, user *llm.UserConte
 	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) {
@@ -45,36 +39,37 @@ func (k *Kimi) IsLogin(ctx context.Context, page *rod.Page) (bool, error) {
 	return true, nil
 }
 
-func (k *Kimi) Login(ctx context.Context, page *rod.Page) error {
+func (k *Kimi) Login(ctx context.Context, page *rod.Page) (*llm.UserContext, error) {
 	loginEle, err := page.Element("div[data-testid='msh-sidebar-user']")
+	userContext := llm.NewUserContext()
 	if err != nil {
-		return err
+		return userContext, err
 	}
 	err = loginEle.Click(proto.InputMouseButtonLeft, 1)
 	if err != nil {
-		return err
+		return userContext, err
 	}
 	qrcodePreEle, err := page.ElementR("div p", "微信扫码登录")
 	if err != nil {
-		return err
+		return userContext, err
 	}
 	qrcodeBox, err := qrcodePreEle.Parent()
 	if err != nil {
-		return err
+		return userContext, err
 	}
 	screenshot, err := qrcodeBox.Screenshot(proto.PageCaptureScreenshotFormatPng, 10)
 	if err != nil {
-		return err
+		return userContext, err
 	}
 
 	err = k.loginEvent(ctx, screenshot)
 	if err != nil {
-		return err
+		return userContext, err
 	}
 
 	_, err = page.Element("div[data-testid='msh-header-user-avatar']")
 	if err != nil {
-		return err
+		return userContext, err
 	}
 
 	localStorage := page.MustEval(`k => Object.assign({}, window.localStorage)`).Map()
@@ -82,12 +77,12 @@ func (k *Kimi) Login(ctx context.Context, page *rod.Page) error {
 	for k, v := range localStorage {
 		localStorageMap[k] = v.String()
 	}
-	k.userContext = llm.NewUserContext().WithLocalStorage(k.GetName(), localStorageMap).WithCookies(k.GetName(), proto.CookiesToParams(page.MustCookies()))
-	err = k.loginCallback(ctx, k.userContext)
+	userContext = llm.NewUserContext().WithLocalStorage(k.GetName(), localStorageMap).WithCookies(k.GetName(), proto.CookiesToParams(page.MustCookies()))
+	err = k.loginCallback(ctx, userContext)
 	if err != nil {
-		return err
+		return userContext, err
 	}
-	return nil
+	return userContext, nil
 }
 
 func (k *Kimi) GetName() string {
@@ -121,8 +116,8 @@ func (k *Kimi) NewChat(ctx context.Context) (llm.Chat, error) {
 	if err != nil {
 		return nil, err
 	}
-	if k.userContext != nil {
-		localStorage, ok := k.userContext.LocalStorage[k.GetName()]
+	if c.userContext != nil {
+		localStorage, ok := c.userContext.LocalStorage[k.GetName()]
 		if ok {
 			for key, val := range localStorage {
 				c.Page.MustEval(fmt.Sprintf("() => {window.localStorage.setItem('%s','%s')}", key, val))

llm.go

@@ -6,6 +6,7 @@ import (
 
 type Chat interface {
 	Ask(ctx context.Context, text string) (string, error)
+	WithUserContext(user *UserContext) Chat
 	Close(ctx context.Context) error
 }
 

user.go

@@ -10,10 +10,9 @@ import (
 
 type UserBrowserAgent interface {
 	IsLogin(ctx context.Context, page *rod.Page) (bool, error)
-	Login(ctx context.Context, page *rod.Page) 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
-	WithUserContext(user *UserContext) UserBrowserAgent
 }
 
 type UserContext struct {