import edward

llm/kimi.py

@@ -0,0 +1,55 @@
+from DrissionPage import Chromium, ChromiumOptions
+from DrissionPage.errors import ElementNotFoundError
+
+
+class Kimi:
+    def __init__(self, base_timeout=30):
+        self.base_timeout = base_timeout
+        """Initialize the Chromium browser with specified timeout settings."""
+        # co = ChromiumOptions().auto_port()
+        co = ChromiumOptions()
+        co.set_timeouts(base=self.base_timeout)
+        self.browser = Chromium(addr_or_opts=co)
+        self.tab = self.browser.latest_tab
+        self.tab.get('https://kimi.moonshot.cn/')
+        self.first_time = True
+
+    def chat(self, input_message, wait_time=10):
+        """Send a message through the chat interface and retrieve the response."""
+        try:
+            chat_input = self.tab.ele('.editorContentEditable___FZJd9')
+            chat_input.wait.enabled()
+            chat_input.input(input_message)
+            self.tab.wait(3)
+            self.tab.ele('#send-button').wait.enabled()
+            self.tab.ele('#send-button').click()
+            self.tab.wait(wait_time)
+            chat_input.input(' ')
+            self.tab.ele('#send-button').wait.enabled()
+            output = self.tab.ele('#^chat-markdown').text
+        except ElementNotFoundError:
+            output = ''
+        return output
+
+    def generate(self, input_message, wait_time=10):
+        if not self.first_time:
+            new_session_btn = self.tab.ele('.myAgentTool___Y1_mC')
+            new_session_btn.wait.enabled()
+            new_session_btn.click()
+            self.tab.wait(3)
+        self.first_time = False
+        return self.chat(input_message, wait_time)
+
+    def search(self, input_message, wait_time=10):
+        return self.chat(input_message, wait_time)
+
+    def quit(self):
+        self.browser.quit()
+
+if __name__ == "__main__":
+    kimi = Kimi()
+    message = '今日新闻'
+    response = kimi.search(message)
+    kimi.quit()
+
+    print(response)