Spaces:
Runtime error
Runtime error
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 |