3
3
4
4
from revChatGPT .V1 import Chatbot
5
5
6
+ from aishell .exceptions import UnauthorizedAccessError
6
7
from aishell .utils import make_executable_command
7
8
8
9
from .query_client import QueryClient
9
10
10
11
11
12
class ReverseEngineeredChatGPTClient (QueryClient ):
12
- access_key : str
13
+ config : dict [ str , str ] = {}
13
14
14
15
def __init__ (
15
16
self ,
16
- chatgpt_access_key : Optional [str ] = None ,
17
+ access_token : Optional [str ] = None ,
18
+ session_token : Optional [str ] = None ,
17
19
):
18
- super ().__init__ ()
19
- CHATGPT_ACCESS_KEY = os .environ .get ('CHATGPT_ACCESS_KEY' )
20
-
21
- if chatgpt_access_key is not None :
22
- self .access_key = chatgpt_access_key
23
- elif CHATGPT_ACCESS_KEY is not None :
24
- self .access_key = CHATGPT_ACCESS_KEY
20
+ CHATGPT_ACCESS_TOKEN = os .environ .get ('CHATGPT_ACCESS_TOKEN' , access_token )
21
+ CHATGPT_SESSION_TOKEN = os .environ .get ('CHATGPT_SESSION_TOKEN' , session_token )
22
+ if CHATGPT_ACCESS_TOKEN is not None :
23
+ self .config ['access_token' ] = CHATGPT_ACCESS_TOKEN
24
+ elif CHATGPT_SESSION_TOKEN is not None :
25
+ self .config ['session_token' ] = CHATGPT_SESSION_TOKEN
25
26
else :
26
- raise Exception ( 'access_key should not be none ' )
27
+ raise UnauthorizedAccessError ( 'No access token or session token provided. ' )
27
28
28
29
def _construct_prompt (self , text : str ) -> str :
29
30
return f'''You are now a translater from human language to { os .uname ()[0 ]} shell command.
@@ -32,7 +33,7 @@ def _construct_prompt(self, text: str) -> str:
32
33
33
34
def query (self , prompt : str ) -> str :
34
35
prompt = self ._construct_prompt (prompt )
35
- chatbot = Chatbot (config = { 'access_token' : self .access_key } )
36
+ chatbot = Chatbot (config = self .config )
36
37
37
38
response_text = ''
38
39
for data in chatbot .ask (prompt ):
0 commit comments