chat.go 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. package kimi
  2. import (
  3. "context"
  4. "errors"
  5. "git.ttmylife.com/zeuszhao/llm"
  6. "github.com/go-rod/rod"
  7. "github.com/go-rod/rod/lib/proto"
  8. "sync"
  9. "time"
  10. )
  11. type Chat struct {
  12. ai llm.AI
  13. Page *rod.Page
  14. qaSingle chan bool
  15. sync.Mutex
  16. }
  17. func (c *Chat) Ask(ctx context.Context, text string) (string, error) {
  18. c.TryLock()
  19. defer c.Unlock()
  20. ok, err := c.ai.IsLogin(ctx, c.Page)
  21. if err != nil {
  22. return "", err
  23. }
  24. if !ok {
  25. err = c.ai.Login(ctx, c.Page)
  26. if err != nil {
  27. return "", err
  28. }
  29. }
  30. inputEle, err := c.Page.Element("#msh-chateditor")
  31. if err != nil {
  32. return "", err
  33. }
  34. err = inputEle.Input(text)
  35. if err != nil {
  36. return "", err
  37. }
  38. sendBtnEle, err := c.Page.Element("#send-button")
  39. if err != nil {
  40. return "", errors.Join(err, errors.New("send btn not fount"))
  41. }
  42. // 等待按钮可用
  43. err = sendBtnEle.WaitWritable()
  44. if err != nil {
  45. return "", errors.Join(err, errors.New("send btn not writable"))
  46. }
  47. err = sendBtnEle.WaitStable(1 * time.Second)
  48. if err != nil {
  49. return "", errors.Join(err, errors.New("send btn not stable"))
  50. }
  51. err = sendBtnEle.Click(proto.InputMouseButtonLeft, 1)
  52. if err != nil {
  53. return "", err
  54. }
  55. select {
  56. case <-c.qaSingle:
  57. case <-ctx.Done():
  58. return "", ctx.Err()
  59. case <-time.After(60 * time.Second):
  60. return "", errors.New("超时")
  61. }
  62. qaEleOk, err := c.Page.Element("button[data-testid='msh-chat-segment-reAnswer']")
  63. if err != nil {
  64. return "", err
  65. }
  66. err = qaEleOk.WaitVisible()
  67. if err != nil {
  68. return "", err
  69. }
  70. qaEle, err := c.Page.Elements(".pop-content")
  71. if err != nil {
  72. return "", err
  73. }
  74. return qaEle.Last().Text()
  75. }
  76. func (c *Chat) Close(ctx context.Context) error {
  77. return c.Page.Close()
  78. }