zeuszhao 3 mēneši atpakaļ
vecāks
revīzija
a3647d65d1
3 mainītis faili ar 11 papildinājumiem un 4 dzēšanām
  1. 8 2
      kimi/kimi.go
  2. 2 1
      llm.go
  3. 1 1
      user.go

+ 8 - 2
kimi/kimi.go

@@ -12,6 +12,7 @@ import (
 )
 
 type Kimi struct {
+	userContext   *llm.UserContext
 	browser       *rod.Browser
 	launcher      *launcher.Launcher
 	loginEvent    func(ctx context.Context, bytes []byte) error
@@ -28,6 +29,11 @@ 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) {
@@ -87,7 +93,7 @@ func (k *Kimi) GetName() string {
 	return "kimi"
 }
 
-func (k *Kimi) NewChat(ctx context.Context, user *llm.UserContext) (llm.Chat, error) {
+func (k *Kimi) NewChat(ctx context.Context) (llm.Chat, error) {
 	c := &Chat{
 		ai:       k,
 		Page:     k.browser.Context(ctx).MustIncognito().MustPage(),
@@ -114,7 +120,7 @@ func (k *Kimi) NewChat(ctx context.Context, user *llm.UserContext) (llm.Chat, er
 	if err != nil {
 		return nil, err
 	}
-	localStorage, ok := user.LocalStorage[k.GetName()]
+	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))

+ 2 - 1
llm.go

@@ -12,6 +12,7 @@ type Chat interface {
 type AI interface {
 	UserBrowserAgent
 	Init(ctx context.Context) (AI, error)
-	NewChat(ctx context.Context, user *UserContext) (Chat, error)
+	NewChat(ctx context.Context) (Chat, error)
 	Close(ctx context.Context) error
+	GetName() string
 }

+ 1 - 1
user.go

@@ -13,7 +13,7 @@ type UserBrowserAgent interface {
 	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
+	WithUserContext(user *UserContext) UserBrowserAgent
 }
 
 type UserContext struct {