Make cli to ask copy session token or not

Author: code-yeongyu

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?')