Create pawansapi.py

pawansapi.py

@@ -0,0 +1,107 @@
+import time
+import uuid
+import asyncio
+import socketio
+import datetime
+import json
+import base64
+from loguru import logger
+from socketio.exceptions import ConnectionError, TimeoutError
+
+class PyGPT:
+    def __init__(self, session_token, bypass_node='https://gpt.pawan.krd', db = None):
+        self.ready = False
+        self.socket = socketio.AsyncClient()
+        self.socket.on('connect', self.on_connect)
+        self.socket.on('disconnect', self.on_disconnect)
+        self.session_token = session_token
+        self.auth = None
+        self.expires = datetime.datetime.now()
+        self.pause_token_checks = False
+        self.bypass_node = bypass_node
+        self.db = db
+
+    async def connect(self, retry=0):
+        if retry<=3:
+            try:          #python&version=1.0.2&versionCode=102   nodejs&version=1.0.5&versionCode=105
+                await self.socket.connect(f'{self.bypass_node}/?client=python&version=1.0.2&versionCode=102')
+            except ConnectionError as e:
+                time.sleep(3)
+                print(f"retry {retry} for connection!")
+                await self.connect(retry+1)
+        else:
+            raise ConnectionError 
+
+    async def disconnect(self):
+        await self.socket.disconnect()
+
+    def on_connect(self):
+        logger.info('Connected to server')
+        asyncio.create_task(self.check_tokens())
+
+    def on_disconnect(self):
+        logger.info('Disconnected from server')
+        self.ready = False
+
+    async def check_tokens(self):
+        while True:
+            if self.pause_token_checks:
+                await asyncio.sleep(0.5)
+                continue
+            self.pause_token_checks = True
+            now = datetime.datetime.now()
+            offset = datetime.timedelta(minutes=2)
+            if self.expires < (now - offset) or not self.auth:
+                await self.get_tokens()
+            self.pause_token_checks = False
+            await asyncio.sleep(0.5)
+
+    async def wait_for_ready(self):
+        while not self.ready:
+            await asyncio.sleep(0.025)
+        print('Ready!!')
+
+    async def ask(self, query_data):
+        
+        ask_data = {
+            'prompt': query_data['message'],
+            'parentId': str(query_data.get('parentId', uuid.uuid4())),
+            'conversationId': str(query_data.get('conversationId', None)),
+            'auth': self.auth
+        }
+#         logger.debug(f"pogpt query_data: prompt-{ask_data['prompt']}, parentId-{ask_data['parentId']}, conversationId-{ask_data['conversationId']}")
+        logger.debug(f"pogpt query_data: prompt-{ask_data}")
+            
+        data = await self.socket.call('askQuestion', ask_data, timeout=150)
+
+        if 'error' in data:
+            print(f'Error: {data}')
+            return {'status': 404, 'reply': data["error"]}
+        try:
+            query_data['parentId'] = data['messageId']
+            query_data['conversationId'] = data['conversationId']
+            query_data['reply'] = data['answer']
+            query_data['status'] = 200
+            return query_data
+        except Exception as e:
+            print(f'Error noerror in data: {data}')
+            return {'status': 404, 'reply': str(data)}
+        
+    def validate_token(self, token):
+        if not token:
+            return False
+        parsed = json.loads(base64.b64decode(f'{token.split(".")[1]}==').decode())
+        print("expire time:", datetime.datetime.fromtimestamp(parsed['exp']))
+        return datetime.datetime.now() <= datetime.datetime.fromtimestamp(parsed['exp'])
+
+    async def get_tokens(self):
+        await asyncio.sleep(1)
+        data = await self.socket.call('getSession', self.session_token, timeout= 60)
+        
+        if 'error' in data:
+            print(f'Error getting session: {data["error"]}')
+        else:
+            self.auth = data['auth']
+            self.expires = datetime.datetime.strptime(data['expires'], '%Y-%m-%dT%H:%M:%S.%fZ')
+            self.session_token = data['sessionToken']
+            self.ready = True