Add auto login feature integrated with settings file 🔥🔥🔥

Fixes #17

Author: code-yeongyu

aishell/cli.py

@@ -1,58 +1,77 @@
 import os
 import sys
-import webbrowser
 
-import pyperclip
 import rich
 import typer
 from rich.console import Console
 from yt_dlp.cookies import SUPPORTED_BROWSERS
 
 from aishell.adapters.openai_cookie_adapter import OpenAICookieAdapter
-from aishell.exceptions import UnauthorizedAccessError
+from aishell.models import RevChatGPTChatbotConfigModel
+from aishell.models.aishell_config_model import AiShellConfigModel
 from aishell.models.language_model import LanguageModel
 from aishell.query_clients import GPT3Client, OfficialChatGPTClient, QueryClient, ReverseEngineeredChatGPTClient
+from aishell.utils import AiShellConfigManager
 
 cli_app = typer.Typer()
 
 
-def _open_chatgpt_browser():
-    CHATGPT_LOGIN_URL = 'https://chat.openai.com/auth/login?next=/chat'
-    webbrowser.open(CHATGPT_LOGIN_URL)
+def config_aishell():
+    rich.print('''Hi! 🙌 I am [bold blue]AiShell[/bold blue], [yellow]your powerful terminal assistant[/yellow] 🔥
+I am here to assist you with configuring AiShell. 💪
 
+Please make sure that you have logged into chat.openai.com on your browser before we continue. 🗝️
 
-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.'
-        )
+''')
+    typer.confirm('Are you ready to proceed? 🚀', abort=True)
+
+    rich.print(f'''Which browser did you use to log in to chat.openai.com?
+
+We support the following browsers: [{SUPPORTED_BROWSERS}]''')
+    browser_name = typer.prompt('Please enter your choice here: ')
+    if browser_name not in SUPPORTED_BROWSERS:
+        rich.print(f'Browser {browser_name} is not supported. Supported browsers are: {SUPPORTED_BROWSERS}')
+        sys.exit(1)
+
+    adapter = OpenAICookieAdapter(browser_name)
+    session_token = adapter.get_openai_session_token()
+    if not session_token:
+        rich.print('Failed to get session token. 😓 Can you check if you are logged in to https://chat.openai.com?')
+        sys.exit(1)
+
+    is_paid = typer.confirm("It's my last question! 🤩 Are you a PLUS user?")
+
+    chatgpt_config = RevChatGPTChatbotConfigModel(session_token=session_token, paid=is_paid)
+    aishell_config = AiShellConfigModel(chatgpt_config=chatgpt_config)
+    config_manager = AiShellConfigManager(config_model=aishell_config)
+    config_manager.save_config()
+
+    rich.print(f'''[green bold]Excellent![/green bold] You are now ready to use [bold blue]AiShell[/bold blue] 🚀
+
+Enjoy your AI powered terminal assistant! 🎉
+
+[dim]To check your settings file, it's at: {config_manager.config_path}[/dim]
+
+''')
+    return config_manager
 
 
 @cli_app.command()
 def ask(question: str, language_model: LanguageModel = LanguageModel.REVERSE_ENGINEERED_CHATGPT):
+    is_config_file_available = AiShellConfigManager.is_config_file_available(AiShellConfigManager.DEFAULT_CONFIG_PATH)
+    config_manager: AiShellConfigManager
+    if is_config_file_available:
+        config_manager = AiShellConfigManager(load_config=True)
+    else:
+        config_manager = config_aishell()
+
     query_client: QueryClient
-    if language_model == LanguageModel.GPT3:
+    if language_model == LanguageModel.REVERSE_ENGINEERED_CHATGPT:
+        query_client = ReverseEngineeredChatGPTClient(config=config_manager.config_model.chatgpt_config)
+    elif language_model == LanguageModel.GPT3:
         query_client = GPT3Client()
     elif language_model == LanguageModel.OFFICIAL_CHATGPT:
         query_client = OfficialChatGPTClient()
-    elif language_model == LanguageModel.REVERSE_ENGINEERED_CHATGPT:
-        try:
-            query_client = ReverseEngineeredChatGPTClient()
-        except UnauthorizedAccessError:
-            print('You are not logged in to OpenAI, attempting to log you in...')
-            _open_chatgpt_browser()
-            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()
-            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)
 

aishell/query_clients/reverse_engineered_chatgpt_client.py

@@ -19,17 +19,21 @@ def revchatgpt_config(self) -> dict[str, Union[str, bool]]:
 
     def __init__(
         self,
+        config: Optional[RevChatGPTChatbotConfigModel] = None,
         access_token: Optional[str] = None,
         session_token: Optional[str] = None,
     ):
-        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:
-            self._config = RevChatGPTChatbotConfigModel(access_token=CHATGPT_ACCESS_TOKEN)
-        elif CHATGPT_SESSION_TOKEN:
-            self._config = RevChatGPTChatbotConfigModel(session_token=CHATGPT_SESSION_TOKEN)
+        if config:
+            self._config = config
         else:
-            raise UnauthorizedAccessError('No access token or session token provided.')
+            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:
+                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 query(self, prompt: str) -> str:
         prompt = self._construct_prompt(prompt)