Ask user to copy session token after retrieving it from chat.openai.com

aishell/cli.py

@@ -1,6 +1,9 @@
 import os
+import sys
 import webbrowser
 
+import pyperclip
+import rich
 import typer
 from rich.console import Console
 from yt_dlp.cookies import SUPPORTED_BROWSERS
@@ -18,6 +21,15 @@ def _open_chatgpt_browser():
     webbrowser.open(CHATGPT_LOGIN_URL)
 
 
+def _ask_user_copy_session_token_to_clipboard(session_token: str) -> None:
+    copy_session_token = typer.confirm('Do you want to copy the session token to your clipboard?')
+    if copy_session_token:
+        pyperclip.copy(session_token)
+        rich.print(
+            'Session token copied to clipboard. [bold]`export CHATGPT_SESSION_TOKEN=<session_token>`[/bold] to set it.'
+        )
+
+
 @cli_app.command()
 def ask(question: str, language_model: LanguageModel = LanguageModel.REVERSE_ENGINEERED_CHATGPT):
     query_client: QueryClient
@@ -34,7 +46,13 @@ def ask(question: str, language_model: LanguageModel = LanguageModel.REVERSE_ENG
             BROWSER_NAME = typer.prompt(f'Which browser did you use to log in? [{SUPPORTED_BROWSERS}]')
             adapter = OpenAICookieAdapter(BROWSER_NAME)
             session_token = adapter.get_openai_session_token()
-            query_client = ReverseEngineeredChatGPTClient(session_token=session_token)
+            if session_token is not None:
+                os.environ['CHATGPT_SESSION_TOKEN'] = session_token
+                _ask_user_copy_session_token_to_clipboard(session_token)
+                ask(question, language_model)
+            else:
+                print('Failed to log in.')
+            sys.exit()
 
     query_client.query(question)
 
@@ -43,9 +61,9 @@ def ask(question: str, language_model: LanguageModel = LanguageModel.REVERSE_ENG
             f'''
 [green] AiShell is thinking of `{question}` ...[/green]
 
-[italic]AiShell is not responsible for any damage caused by the command executed by the user.[/italic]'''.strip(), ):
+[dim]AiShell is not responsible for any damage caused by the command executed by the user.[/dim]'''.strip(), ):
         response = query_client.query(question)
-    console.print(f'[italic]ai$hell: {response}\n')
+    console.print(f'AiShell: {response}\n')
 
     will_execute = typer.confirm('Execute this command?')
 

aishell/models/__init__.py

@@ -1,2 +1,3 @@
 from .language_model import LanguageModel as LanguageModel
 from .open_ai_response_model import OpenAIResponseModel as OpenAIResponseModel
+from .revchatgpt_chatbot_config_model import RevChatGPTChatbotConfigModel as RevChatGPTChatbotConfigModel

aishell/models/open_ai_response_model.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import Optional
 
 from pydantic import BaseModel
 
@@ -16,9 +16,12 @@ class Usage(BaseModel):
         prompt_tokens: int
         total_tokens: int
 
-    choices: Optional[List[Choice]]
+    choices: Optional[list[Choice]]
     created: int
     id: str
     model: str
     object: str
     usage: Usage
+
+    class Config:
+        frozen = True

aishell/models/revchatgpt_chatbot_config_model.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+from pydantic import BaseModel, root_validator
+
+
+class RevChatGPTChatbotConfigModel(BaseModel):
+    email: Optional[str] = None
+    password: Optional[str] = None
+    session_token: Optional[str] = None
+    access_token: Optional[str] = None
+    paid: bool = False
+
+    @root_validator
+    def check_at_least_one_account_info(cls, values: dict[str, Optional[str]]):
+        IS_ACCOUNT_LOGIN = values.get('email') and values.get('password')
+        IS_TOKEN_AUTH = values.get('session_token') or values.get('access_token')
+        if not IS_ACCOUNT_LOGIN and not IS_TOKEN_AUTH:
+            raise ValueError('No information for authentication provided.')
+
+        return values
+
+    class Config:
+        frozen = True

aishell/query_clients/gpt3_client.py

@@ -1,5 +1,5 @@
 import os
-from typing import cast
+from typing import Final, cast
 
 import openai
 
@@ -11,16 +11,9 @@
 
 class GPT3Client(QueryClient):
 
-    def _construct_prompt(self, text: str) -> str:
-        return f'''User: You are now a translater from human language to {os.uname()[0]} shell command.
-        No explanation required, respond with only the raw shell command.
-        What should I type to shell for: {text}, in one line.
-
-        You: '''
-
     def query(self, prompt: str) -> str:
         prompt = self._construct_prompt(prompt)
-        completion: OpenAIResponseModel = cast(  # type: ignore [no-any-unimported]
+        completion: Final[OpenAIResponseModel] = cast(
             OpenAIResponseModel,
             openai.Completion.create(
                 engine='text-davinci-003',
@@ -32,5 +25,12 @@ def query(self, prompt: str) -> str:
         )
         if not completion.choices or len(completion.choices) == 0 or not completion.choices[0].text:
             raise RuntimeError('No response from OpenAI')
-        response_text: str = completion.choices[0].text
+        response_text: Final[str] = completion.choices[0].text
         return make_executable_command(response_text)
+
+    def _construct_prompt(self, text: str) -> str:
+        return f'''User: You are now a translater from human language to {os.uname()[0]} shell command.
+        No explanation required, respond with only the raw shell command.
+        What should I type to shell for: {text}, in one line.
+
+        You: '''

aishell/query_clients/official_chatgpt_client.py

@@ -1,5 +1,5 @@
 import os
-from typing import Final, Optional
+from typing import Optional
 
 from revChatGPT.V3 import Chatbot
 
@@ -10,28 +10,27 @@
 
 
 class OfficialChatGPTClient(QueryClient):
-    openai_api_key: str
 
     def __init__(
         self,
         openai_api_key: Optional[str] = None,
     ):
         super().__init__()
-        OPENAI_API_KEY: Final[Optional[str]] = os.environ.get('OPENAI_API_KEY', openai_api_key)
+        OPENAI_API_KEY: Optional[str] = os.environ.get('OPENAI_API_KEY', openai_api_key)
         if OPENAI_API_KEY is None:
             raise UnauthorizedAccessError('OPENAI_API_KEY should not be none')
 
-        self.openai_api_key = OPENAI_API_KEY
-
-    def _construct_prompt(self, text: str) -> str:
-        return f'''You are now a translater from human language to {os.uname()[0]} shell command.
-        No explanation required, respond with only the raw shell command.
-        What should I type to shell for: {text}, in one line.'''
+        self.OPENAI_API_KEY = OPENAI_API_KEY
 
     def query(self, prompt: str) -> str:
-        prompt = self._construct_prompt(prompt)
+        chatbot = Chatbot(api_key=self.OPENAI_API_KEY)
 
-        chatbot = Chatbot(api_key=self.openai_api_key)
+        prompt = self._construct_prompt(prompt)
         response_text = chatbot.ask(prompt)
+        executable_command = make_executable_command(response_text)
+        return executable_command
 
-        return make_executable_command(response_text)
+    def _construct_prompt(self, text: str) -> str:
+        return f'''You are now a translater from human language to {os.uname()[0]} shell command.
+        No explanation required, respond with only the raw shell command.
+        What should I type to shell for: {text}, in one line.'''

aishell/query_clients/reverse_engineered_chatgpt_client.py

@@ -1,16 +1,21 @@
 import os
-from typing import Optional, cast
+from typing import Optional, Union, cast
 
 from revChatGPT.V1 import Chatbot
 
 from aishell.exceptions import UnauthorizedAccessError
+from aishell.models import RevChatGPTChatbotConfigModel
 from aishell.utils import make_executable_command
 
 from .query_client import QueryClient
 
 
 class ReverseEngineeredChatGPTClient(QueryClient):
-    config: dict[str, str] = {}
+    _config: RevChatGPTChatbotConfigModel
+
+    @property
+    def revchatgpt_config(self) -> dict[str, Union[str, bool]]:
+        return self._config.dict(exclude_none=True)
 
     def __init__(
         self,
@@ -19,21 +24,17 @@ def __init__(
     ):
         CHATGPT_ACCESS_TOKEN = os.environ.get('CHATGPT_ACCESS_TOKEN', access_token)
         CHATGPT_SESSION_TOKEN = os.environ.get('CHATGPT_SESSION_TOKEN', session_token)
-        if CHATGPT_ACCESS_TOKEN is not None:
-            self.config['access_token'] = CHATGPT_ACCESS_TOKEN
-        elif CHATGPT_SESSION_TOKEN is not None:
-            self.config['session_token'] = CHATGPT_SESSION_TOKEN
+        if CHATGPT_ACCESS_TOKEN:
+            self._config = RevChatGPTChatbotConfigModel(access_token=CHATGPT_ACCESS_TOKEN)
+        elif CHATGPT_SESSION_TOKEN:
+            self._config = RevChatGPTChatbotConfigModel(session_token=CHATGPT_SESSION_TOKEN)
         else:
             raise UnauthorizedAccessError('No access token or session token provided.')
 
-    def _construct_prompt(self, text: str) -> str:
-        return f'''You are now a translater from human language to {os.uname()[0]} shell command.
-        No explanation required, respond with only the raw shell command.
-        What should I type to shell for: {text}, in one line.'''
-
     def query(self, prompt: str) -> str:
         prompt = self._construct_prompt(prompt)
-        chatbot = Chatbot(config=self.config)
+        chatbot = Chatbot(config=self.revchatgpt_config)  #  pyright: ignore [reportGeneralTypeIssues]
+        # ignore for wrong type hint of revchatgpt
 
         response_text = ''
         for data in chatbot.ask(prompt):
@@ -42,3 +43,8 @@ def query(self, prompt: str) -> str:
         response_text = make_executable_command(cast(str, response_text))
 
         return response_text
+
+    def _construct_prompt(self, text: str) -> str:
+        return f'''You are now a translater from human language to {os.uname()[0]} shell command.
+        No explanation required, respond with only the raw shell command.
+        What should I type to shell for: {text}, in one line.'''

aishell/utils/str_enum.py

@@ -4,7 +4,12 @@
 
 class StrEnum(str, Enum):
 
-    def _generate_next_value_(name: str, start: int, count: int, last_values: list[Any]):  # type: ignore
+    def _generate_next_value_(  # pyright: ignore [reportIncompatibleMethodOverride], for pyright's bug
+        name: str,
+        start: int,
+        count: int,
+        last_values: list[Any],
+    ):
         return name.lower()
 
     def __repr__(self):

pyproject.toml

@@ -47,6 +47,7 @@ openai = "^0.26.5"
 pydantic = "^1.10.4"
 pyright = "^1.1.294"
 yt-dlp = "^2023.2.17"
+pyperclip = "^1.8.2"
 
 [tool.poetry.scripts]
 aishell = "aishell:main"