Yusin commited on
Commit
a53b4a8
1 Parent(s): 93641fd

Create pygpt.py

Browse files
Files changed (1) hide show
  1. pygpt.py +111 -0
pygpt.py ADDED
@@ -0,0 +1,111 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import uuid
2
+ import asyncio
3
+ import socketio
4
+ import datetime
5
+ import json
6
+ import base64
7
+
8
+ class PyGPT:
9
+ def __init__(self, session_token, bypass_node='https://gpt.pawan.krd'):
10
+ self.ready = False
11
+ self.socket = socketio.AsyncClient()
12
+ self.socket.on('connect', self.on_connect)
13
+ self.socket.on('disconnect', self.on_disconnect)
14
+ self.session_token = session_token
15
+ self.conversations = []
16
+ self.auth = None
17
+ self.expires = datetime.datetime.now()
18
+ self.pause_token_checks = False
19
+ self.bypass_node = bypass_node
20
+ asyncio.create_task(self.cleanup_conversations())
21
+
22
+ async def connect(self):
23
+ await self.socket.connect(self.bypass_node)
24
+
25
+ async def disconnect(self):
26
+ await self.socket.disconnect()
27
+
28
+ def on_connect(self):
29
+ print('Connected to server')
30
+ asyncio.create_task(self.check_tokens())
31
+
32
+ def on_disconnect(self):
33
+ print('Disconnected from server')
34
+ self.ready = False
35
+
36
+ async def check_tokens(self):
37
+ while True:
38
+ if self.pause_token_checks:
39
+ await asyncio.sleep(0.5)
40
+ continue
41
+ self.pause_token_checks = True
42
+ now = datetime.datetime.now()
43
+ offset = datetime.timedelta(minutes=2)
44
+ if self.expires < (now - offset) or not self.auth:
45
+ await self.get_tokens()
46
+ self.pause_token_checks = False
47
+ await asyncio.sleep(0.5)
48
+
49
+ async def cleanup_conversations(self):
50
+ while True:
51
+ await asyncio.sleep(60)
52
+ now = datetime.datetime.now()
53
+ self.conversations = [c for c in self.conversations if now - c['last_active'] < datetime.timedelta(minutes=2)]
54
+
55
+ def add_conversation(self, id):
56
+ conversation = {
57
+ 'id': id,
58
+ 'conversation_id': None,
59
+ 'parent_id': uuid.uuid4(),
60
+ 'last_active': datetime.datetime.now()
61
+ }
62
+ self.conversations.append(conversation)
63
+ return conversation
64
+
65
+ def get_conversation_by_id(self, id):
66
+ conversation = next((c for c in self.conversations if c['id'] == id), None)
67
+ if conversation is None:
68
+ conversation = self.add_conversation(id)
69
+ else:
70
+ conversation['last_active'] = datetime.datetime.now()
71
+ return conversation
72
+
73
+ async def wait_for_ready(self):
74
+ while not self.ready:
75
+ await asyncio.sleep(0.025)
76
+ print('Ready!!')
77
+
78
+ async def ask(self, prompt, id='default'):
79
+ if not self.auth or not self.validate_token(self.auth):
80
+ await self.get_tokens()
81
+ conversation = self.get_conversation_by_id(id)
82
+ data = await self.socket.call('askQuestion', {
83
+ 'prompt': prompt,
84
+ 'parentId': str(conversation['parent_id']),
85
+ 'conversationId': str(conversation['conversation_id']),
86
+ 'auth': self.auth
87
+ })
88
+
89
+ if 'error' in data:
90
+ print(f'Error: {data["error"]}')
91
+ conversation['parent_id'] = data['messageId']
92
+ conversation['conversation_id'] = data['conversationId']
93
+ return data['answer']
94
+
95
+ def validate_token(self, token):
96
+ if not token:
97
+ return False
98
+ parsed = json.loads(base64.b64decode(f'{token.split(".")[1]}==').decode())
99
+ return datetime.datetime.now() <= datetime.datetime.fromtimestamp(parsed['exp'])
100
+
101
+ async def get_tokens(self):
102
+ await asyncio.sleep(1)
103
+ data = await self.socket.call('getSession', self.session_token)
104
+
105
+ if 'error' in data:
106
+ print(f'Error getting session: {data["error"]}')
107
+ else:
108
+ self.auth = data['auth']
109
+ self.expires = datetime.datetime.strptime(data['expires'], '%Y-%m-%dT%H:%M:%S.%fZ')
110
+ self.session_token = data['sessionToken']
111
+ self.ready = True