6 files changed, 246 insertions, 41 deletions

diff --git a/poe/__init__.py b/poe/__init__.py
index 3f75ba14..d1cb8ae5 100644
--- a/poe/__init__.py
+++ b/poe/__init__.py
@@ -1,6 +1,6 @@
 from poe.api    import Client as PoeClient
 from poe.mail   import Mail
-from tls_client import Session
+from requests   import Session
 from re         import search, findall
 from json       import loads
 from time       import sleep, time
@@ -48,11 +48,10 @@ class PoeResponse:
     def json(self) -> dict:
         return self.response_dict
 
-
 class Account:
     def create(proxy: None or str = None, logging: bool = False):
         
-        client       = Session(client_identifier = "chrome110")
+        client       = Session()
         client.proxies = {
             'http': f'http://{proxy}',
             'https': f'http://{proxy}'} if proxy else None
diff --git a/poe/api.py b/poe/api.py
index baaa1338..e4c5d166 100644
--- a/poe/api.py
+++ b/poe/api.py
@@ -22,6 +22,7 @@ import logging
 import time
 import queue
 import threading
+import traceback
 import websocket
 from pathlib import Path
 from urllib.parse import urlparse
@@ -124,6 +125,15 @@ class Client:
 
         return next_data
 
+    def get_bot(self, display_name):
+        url = f'https://poe.com/_next/data/{self.next_data["buildId"]}/{display_name}.json'
+        logger.info("Downloading "+url)
+
+        r = request_with_retries(self.session.get, url)
+
+        chat_data = r.json()["pageProps"]["payload"]["chatOfBotDisplayName"]
+        return chat_data
+
     def get_bots(self):
         viewer = self.next_data["props"]["pageProps"]["payload"]["viewer"]
         if not "availableBots" in viewer:
@@ -132,13 +142,7 @@ class Client:
 
         bots = {}
         for bot in bot_list:
-            url = f'https://poe.com/_next/data/{self.next_data["buildId"]}/{bot["displayName"].lower()}.json'
-            logger.info("Downloading "+url)
-
-            r = request_with_retries(self.session.get, url)
-
-            chat_data = r.json()[
-                "pageProps"]["payload"]["chatOfBotDisplayName"]
+            chat_data = self.get_bot(bot["displayName"].lower())
             bots[chat_data["defaultBotObject"]["nickname"]] = chat_data
 
         return bots
@@ -165,11 +169,8 @@ class Client:
         return f'wss://{self.ws_domain}.tch.{channel["baseHost"]}/up/{channel["boxName"]}/updates'+query
 
     def send_query(self, query_name, variables):
-        # print(f'send_query: {query_name} {variables}')
-        
         for i in range(20):
             payload = generate_payload(query_name, variables)
-            # print(f'query_payload: {query_name} {variables}')
             r = request_with_retries(
                 self.session.post, self.gql_url, json=payload, headers=self.gql_headers)
             data = r.json()
@@ -216,7 +217,8 @@ class Client:
             header={"User-Agent": user_agent},
             on_message=self.on_message,
             on_open=self.on_ws_connect,
-            on_error=self.on_ws_error
+            on_error=self.on_ws_error,
+            on_close=self.on_ws_close
         )
         t = threading.Thread(target=self.ws_run_thread, daemon=True)
         t.start()
@@ -231,27 +233,44 @@ class Client:
     def on_ws_connect(self, ws):
         self.ws_connected = True
 
+    def on_ws_close(self, ws, close_status_code):
+        self.ws_connected = False
+        logger.warn(f"Websocket closed with status {close_status_code}")
+
     def on_ws_error(self, ws, error):
-        logger.warn(f"Websocket returned error: {error}")
         self.disconnect_ws()
         self.connect_ws()
 
     def on_message(self, ws, msg):
-        data = json.loads(msg)
-        message = json.loads(data["messages"][0])[
-            "payload"]["data"]["messageAdded"]
-
-        copied_dict = self.active_messages.copy()
-        for key, value in copied_dict.items():
-            # add the message to the appropriate queue
-            if value == message["messageId"] and key in self.message_queues:
-                self.message_queues[key].put(message)
+        try:
+            data = json.loads(msg)
+
+            if not "messages" in data:
                 return
 
-            # indicate that the response id is tied to the human message id
-            elif key != "pending" and value == None and message["state"] != "complete":
-                self.active_messages[key] = message["messageId"]
-                self.message_queues[key].put(message)
+            for message_str in data["messages"]:
+                message_data = json.loads(message_str)
+                if message_data["message_type"] != "subscriptionUpdate":
+                    continue
+                message = message_data["payload"]["data"]["messageAdded"]
+
+                copied_dict = self.active_messages.copy()
+                for key, value in copied_dict.items():
+                    # add the message to the appropriate queue
+                    if value == message["messageId"] and key in self.message_queues:
+                        self.message_queues[key].put(message)
+                        return
+
+                    # indicate that the response id is tied to the human message id
+                    elif key != "pending" and value == None and message["state"] != "complete":
+                        self.active_messages[key] = message["messageId"]
+                        self.message_queues[key].put(message)
+                        return
+
+        except Exception:
+            logger.error(traceback.format_exc())
+            self.disconnect_ws()
+            self.connect_ws()
 
     def send_message(self, chatbot, message, with_chat_break=False, timeout=20):
         # if there is another active message, wait until it has finished sending
@@ -262,8 +281,11 @@ class Client:
         self.active_messages["pending"] = None
 
         logger.info(f"Sending message to {chatbot}: {message}")
-
-        message_data = self.send_query("AddHumanMessageMutation", {
+        # reconnect websocket
+        if not self.ws_connected:
+            self.disconnect_ws()
+            self.connect_ws()
+        message_data = self.send_query("SendMessageMutation", {
             "bot": chatbot,
             "query": message,
             "chatId": self.bots[chatbot]["chatId"],
@@ -272,11 +294,11 @@ class Client:
         })
         del self.active_messages["pending"]
 
-        if not message_data["data"]["messageCreateWithStatus"]["messageLimit"]["canSend"]:
+        if not message_data["data"]["messageEdgeCreate"]["message"]:
             raise RuntimeError(f"Daily limit reached for {chatbot}.")
         try:
-            human_message = message_data["data"]["messageCreateWithStatus"]
-            human_message_id = human_message["message"]["messageId"]
+            human_message = message_data["data"]["messageEdgeCreate"]["message"]
+            human_message_id = human_message["node"]["messageId"]
         except TypeError:
             raise RuntimeError(
                 f"An unknown error occured. Raw response data: {message_data}")
@@ -313,4 +335,67 @@ class Client:
         del self.active_messages[human_message_id]
         del self.message_queues[human_message_id]
 
-load_queries()
+    def send_chat_break(self, chatbot):
+        logger.info(f"Sending chat break to {chatbot}")
+        result = self.send_query("AddMessageBreakMutation", {
+            "chatId": self.bots[chatbot]["chatId"]
+        })
+        return result["data"]["messageBreakCreate"]["message"]
+
+    def get_message_history(self, chatbot, count=25, cursor=None):
+        logger.info(f"Downloading {count} messages from {chatbot}")
+
+        if cursor == None:
+            chat_data = self.get_bot(self.bot_names[chatbot])
+            if not chat_data["messagesConnection"]["edges"]:
+                return []
+            cursor = chat_data["messagesConnection"]["edges"][-1]["cursor"]
+
+        cursor = str(cursor)
+        if count > 50:
+            messages = self.get_message_history(
+                chatbot, count=50, cursor=cursor)
+            while count > 0:
+                new_cursor = messages[0]["cursor"]
+                new_messages = self.get_message_history(
+                    chatbot, min(50, count), cursor=new_cursor)
+                messages = new_messages + messages
+                count -= 50
+            return messages
+
+        result = self.send_query("ChatListPaginationQuery", {
+            "count": count,
+            "cursor": cursor,
+            "id": self.bots[chatbot]["id"]
+        })
+        return result["data"]["node"]["messagesConnection"]["edges"]
+
+    def delete_message(self, message_ids):
+        logger.info(f"Deleting messages: {message_ids}")
+        if not type(message_ids) is list:
+            message_ids = [int(message_ids)]
+
+        result = self.send_query("DeleteMessageMutation", {
+            "messageIds": message_ids
+        })
+
+    def purge_conversation(self, chatbot, count=-1):
+        logger.info(f"Purging messages from {chatbot}")
+        last_messages = self.get_message_history(chatbot, count=50)[::-1]
+        while last_messages:
+            message_ids = []
+            for message in last_messages:
+                if count == 0:
+                    break
+                count -= 1
+                message_ids.append(message["node"]["messageId"])
+
+            self.delete_message(message_ids)
+
+            if count == 0:
+                return
+            last_messages = self.get_message_history(chatbot, count=50)[::-1]
+        logger.info(f"No more messages left to delete.")
+
+
+load_queries()
\ No newline at end of file
diff --git a/poe/cookies.txt b/poe/cookies.txt
index 26850464..0cbd6ca9 100644
--- a/poe/cookies.txt
+++ b/poe/cookies.txt
@@ -1 +1,7 @@
 SmPiNXZI9hBTuf3viz74PA==
+zw7RoKQfeEehiaelYMRWeA==
+NEttgJ_rRQdO05Tppx6hFw==
+3OnmC0r9njYdNWhWszdQJg==
+8hZKR7MxwUTEHvO45TEViw==
+Eea6BqK0AmosTKzoI3AAow==
+pUEbtxobN_QUSpLIR8RGww==
diff --git a/poe/graphql/ChatListPaginationQuery.graphql b/poe/graphql/ChatListPaginationQuery.graphql
index 0bbaf568..6d9ae884 100644
--- a/poe/graphql/ChatListPaginationQuery.graphql
+++ b/poe/graphql/ChatListPaginationQuery.graphql
@@ -11,6 +11,12 @@ query ChatListPaginationQuery(
 }
 
 fragment BotImage_bot on Bot {
+  displayName
+  ...botHelpers_useDeletion_bot
+  ...BotImage_useProfileImage_bot
+}
+
+fragment BotImage_useProfileImage_bot on Bot {
   image {
     __typename
     ... on LocalBotImage {
@@ -20,7 +26,7 @@ fragment BotImage_bot on Bot {
       url
     }
   }
-  displayName
+  ...botHelpers_useDeletion_bot
 }
 
 fragment ChatMessageDownvotedButton_message on Message {
@@ -33,7 +39,7 @@ fragment ChatMessageDropdownMenu_message on Message {
   messageId
   vote
   text
-  linkifiedText
+  author
   ...chatHelpers_isBotMessage
 }
 
@@ -54,6 +60,9 @@ fragment ChatMessageInputView_chat on Chat {
       dailyBalance
       shouldShowRemainingMessageCount
     }
+    hasClearContext
+    isDown
+    ...botHelpers_useDeletion_bot
     id
   }
   shouldShowDisclaimer
@@ -88,6 +97,10 @@ fragment ChatMessageSuggestedReplies_SuggestedReplyButton_message on Message {
 fragment ChatMessageSuggestedReplies_chat on Chat {
   ...ChatWelcomeView_chat
   ...ChatMessageSuggestedReplies_SuggestedReplyButton_chat
+  defaultBotObject {
+    hasWelcomeTopics
+    id
+  }
 }
 
 fragment ChatMessageSuggestedReplies_message on Message {
@@ -97,10 +110,13 @@ fragment ChatMessageSuggestedReplies_message on Message {
 
 fragment ChatMessage_chat on Chat {
   defaultBotObject {
-    ...ChatPageDisclaimer_bot
+    hasWelcomeTopics
+    hasSuggestedReplies
+    disclaimerText
     messageLimit {
       ...ChatPageRateLimitedBanner_messageLimit
     }
+    ...ChatPageDisclaimer_bot
     id
   }
   ...ChatMessageSuggestedReplies_chat
@@ -114,6 +130,7 @@ fragment ChatMessage_message on Message {
   author
   linkifiedText
   state
+  contentType
   ...ChatMessageSuggestedReplies_message
   ...ChatMessageFeedbackButtons_message
   ...ChatMessageOverflowButton_message
@@ -122,12 +139,15 @@ fragment ChatMessage_message on Message {
   ...chatHelpers_isChatBreak
   ...chatHelpers_useTimeoutLevel
   ...MarkdownLinkInner_message
+  ...IdAnnotation_node
 }
 
 fragment ChatMessagesView_chat on Chat {
   ...ChatMessage_chat
   ...ChatWelcomeView_chat
+  ...IdAnnotation_node
   defaultBotObject {
+    hasWelcomeTopics
     messageLimit {
       ...ChatPageRateLimitedBanner_messageLimit
     }
@@ -152,23 +172,42 @@ fragment ChatPageDeleteFooter_chat on Chat {
 }
 
 fragment ChatPageDisclaimer_bot on Bot {
-  disclaimer
+  disclaimerText
+}
+
+fragment ChatPageMainFooter_chat on Chat {
+  defaultBotObject {
+    ...ChatPageMainFooter_useAccessMessage_bot
+    id
+  }
+  ...ChatMessageInputView_chat
+  ...ChatPageShareFooter_chat
+  ...ChatPageDeleteFooter_chat
+}
+
+fragment ChatPageMainFooter_edges on MessageEdge {
+  ...ChatMessageInputView_edges
+}
+
+fragment ChatPageMainFooter_useAccessMessage_bot on Bot {
+  ...botHelpers_useDeletion_bot
+  ...botHelpers_useViewerCanAccessPrivateBot
 }
 
 fragment ChatPageMain_chat_1G22uz on Chat {
   id
   chatId
-  ...ChatMessageInputView_chat
   ...ChatPageShareFooter_chat
   ...ChatPageDeleteFooter_chat
   ...ChatMessagesView_chat
   ...MarkdownLinkInner_chat
   ...chatHelpers_useUpdateStaleChat_chat
   ...ChatSubscriptionPaywallContextWrapper_chat
+  ...ChatPageMainFooter_chat
   messagesConnection(last: $count, before: $cursor) {
     edges {
       ...ChatMessagesView_edges
-      ...ChatMessageInputView_edges
+      ...ChatPageMainFooter_edges
       ...MarkdownLinkInner_edges
       node {
         ...chatHelpers_useUpdateStaleChat_message
@@ -217,6 +256,11 @@ fragment ChatWelcomeView_chat on Chat {
   }
 }
 
+fragment IdAnnotation_node on Node {
+  __isNode: __typename
+  id
+}
+
 fragment MarkdownLinkInner_chat on Chat {
   id
   chatId
@@ -263,6 +307,15 @@ fragment SubscriptionPaywallModal_bot on Bot {
   ...BotImage_bot
 }
 
+fragment botHelpers_useDeletion_bot on Bot {
+  deletionState
+}
+
+fragment botHelpers_useViewerCanAccessPrivateBot on Bot {
+  isPrivateBot
+  viewerIsCreator
+}
+
 fragment chatHelpers_isBotMessage on Message {
   ...chatHelpers_isHumanMessage
   ...chatHelpers_isChatBreak
@@ -292,8 +345,8 @@ fragment chatHelpers_useSendMessage_chat on Chat {
   id
   chatId
   defaultBotObject {
-    nickname
     id
+    nickname
   }
   shouldShowDisclaimer
 }
@@ -303,10 +356,19 @@ fragment chatHelpers_useTimeoutLevel on Message {
   state
   text
   messageId
+  chat {
+    chatId
+    defaultBotNickname
+    id
+  }
 }
 
 fragment chatHelpers_useUpdateStaleChat_chat on Chat {
   chatId
+  defaultBotObject {
+    contextClearWindowSecs
+    id
+  }
   ...chatHelpers_useSendChatBreak_chat
 }
 
diff --git a/poe/graphql/SendMessageMutation.graphql b/poe/graphql/SendMessageMutation.graphql
new file mode 100644
index 00000000..4b0a4383
--- /dev/null
+++ b/poe/graphql/SendMessageMutation.graphql
@@ -0,0 +1,40 @@
+mutation chatHelpers_sendMessageMutation_Mutation(
+  $chatId: BigInt!
+  $bot: String!
+  $query: String!
+  $source: MessageSource
+  $withChatBreak: Boolean!
+) {
+  messageEdgeCreate(chatId: $chatId, bot: $bot, query: $query, source: $source, withChatBreak: $withChatBreak) {
+    chatBreak {
+      cursor
+      node {
+        id
+        messageId
+        text
+        author
+        suggestedReplies
+        creationTime
+        state
+      }
+      id
+    }
+    message {
+      cursor
+      node {
+        id
+        messageId
+        text
+        author
+        suggestedReplies
+        creationTime
+        state
+        chat {
+          shouldShowDisclaimer
+          id
+        }
+      }
+      id
+    }
+  }
+}
diff --git a/testing/poe_test.py b/testing/poe_test.py
new file mode 100644
index 00000000..99c3bea4
--- /dev/null
+++ b/testing/poe_test.py
@@ -0,0 +1,13 @@
+import poe
+from time import sleep
+
+token = poe.Account.create(proxy = 'xtekky:ogingoi2n3g@geo.iproyal.com:12321',logging = True)
+print('token', token)
+
+sleep(2)
+
+for response in poe.StreamingCompletion.create(model  = 'gpt-4',
+    prompt = 'hello world',
+    token  = token):
+    
+    print(response.completion.choices[0].text, end="", flush=True)
\ No newline at end of file