Add reconfigure feature when session_key is not usable

aishell/cli/__init__.py

@@ -0,0 +1,2 @@
+from .aishell import aishell_command as _  # noqa
+from .cli_app import cli_app as cli_app

aishell/cli/aishell.py

@@ -0,0 +1,72 @@
+import os
+import time
+from typing import Optional
+
+import rich
+import typer
+from rich.console import Console
+
+from aishell.models import AiShellConfigModel
+from aishell.models.language_model import LanguageModel
+from aishell.query_clients import GPT3Client, OfficialChatGPTClient, ReverseEngineeredChatGPTClient
+from aishell.utils import AiShellConfigManager
+
+from .cli_app import cli_app
+from .config_aishell import config_aishell
+
+
+@cli_app.command()
+def aishell_command(question: str, language_model: Optional[LanguageModel] = None):
+    config_manager = _get_config_manager()
+    config_manager.config_model.language_model = language_model or config_manager.config_model.language_model
+
+    query_client = _get_query_client(
+        language_model=config_manager.config_model.language_model,
+        config_model=config_manager.config_model,
+    )
+
+    console = Console()
+
+    try:
+        with console.status(f'''
+[green] AiShell is thinking of `{question}` ...[/green]
+
+[dim]AiShell is not responsible for any damage caused by the command executed by the user.[/dim]'''[1:]):
+            start_time = time.time()
+            response = query_client.query(question)
+            end_time = time.time()
+
+        execution_time = end_time - start_time
+        console.print(f'AiShell: {response}\n\n[dim]Took {execution_time:.2f} seconds to think the command.[/dim]')
+
+        will_execute = typer.confirm('Execute this command?')
+        if will_execute:
+            os.system(response)
+    except KeyError:
+        rich.print('It looks like the `session_token` is expired. Please reconfigure AiShell.')
+        typer.confirm('Reconfigure AiShell?', abort=True)
+        config_aishell()
+        aishell_command(question=question, language_model=language_model)
+        typer.Exit()
+
+
+def _get_config_manager():
+    is_config_file_available = AiShellConfigManager.is_config_file_available(AiShellConfigManager.DEFAULT_CONFIG_PATH)
+    if is_config_file_available:
+        return AiShellConfigManager(load_config=True)
+    else:
+        return config_aishell()
+
+
+def _get_query_client(language_model: LanguageModel, config_model: AiShellConfigModel):
+    if language_model == LanguageModel.REVERSE_ENGINEERED_CHATGPT:
+        return ReverseEngineeredChatGPTClient(config=config_model.chatgpt_config)
+
+    if not config_model.openai_api_key:
+        raise RuntimeError('OpenAI API key is not provided. Please provide it in the config file.')
+
+    if language_model == LanguageModel.GPT3:
+        return GPT3Client(openai_api_key=config_model.openai_api_key)
+    if language_model == LanguageModel.OFFICIAL_CHATGPT:
+        return OfficialChatGPTClient(openai_api_key=config_model.openai_api_key)
+    raise NotImplementedError(f'Language model {language_model} is not implemented yet.')

aishell/cli/cli_app.py

@@ -0,0 +1,3 @@
+from typer import Typer
+
+cli_app = Typer()

aishell/cli/config_aishell.py

@@ -0,0 +1,67 @@
+import sys
+
+import rich
+import typer
+from yt_dlp.cookies import SUPPORTED_BROWSERS
+
+from aishell.adapters.openai_cookie_adapter import OpenAICookieAdapter
+from aishell.models import RevChatGPTChatbotConfigModel
+from aishell.models.aishell_config_model import AiShellConfigModel
+from aishell.utils import AiShellConfigManager
+
+
+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. 🗝️
+
+'''[1:])
+    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}]'''[1:])
+    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)
+
+    config_manager = save_config(session_token)
+
+    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]
+
+'''[1:])
+    return config_manager
+
+
+def save_config(session_token: str):
+    is_config_file_available = AiShellConfigManager.is_config_file_available(AiShellConfigManager.DEFAULT_CONFIG_PATH)
+    if is_config_file_available:
+        config_manager = AiShellConfigManager(load_config=True)
+        is_chatgpt_config_available = config_manager.config_model.chatgpt_config is not None
+        if is_chatgpt_config_available:
+            assert config_manager.config_model.chatgpt_config  # for type hinting
+            config_manager.config_model.chatgpt_config.session_token = session_token
+        else:
+            config_manager.config_model.chatgpt_config = RevChatGPTChatbotConfigModel(session_token=session_token)
+    else:
+        chatgpt_config = RevChatGPTChatbotConfigModel(session_token=session_token)
+        aishell_config = AiShellConfigModel(chatgpt_config=chatgpt_config)
+        config_manager = AiShellConfigManager(config_model=aishell_config)
+
+    config_manager.save_config()
+    return config_manager

aishell/cli/test/test_config_aishell.py

@@ -0,0 +1,63 @@
+from unittest.mock import MagicMock, patch
+
+from aishell.cli.config_aishell import save_config
+from aishell.models import AiShellConfigModel, LanguageModel, RevChatGPTChatbotConfigModel
+from aishell.utils import AiShellConfigManager
+
+
+class TestSaveConfig:
+
+    @patch('aishell.cli.config_aishell.AiShellConfigManager', spec=AiShellConfigManager)
+    def test_success_when_config_file_not_available(
+        self,
+        mocked_config_manager_class: MagicMock,
+    ):
+        '''config file 이 없는 경우 테스트 성공해야 한다'''
+        # given
+        mocked_config_manager_class.is_config_file_available.return_value = False
+
+        # when
+        save_config('valid_session_token')
+
+        # then
+        mocked_config_manager_class.return_value.save_config.assert_called_once()
+
+    @patch('aishell.cli.config_aishell.AiShellConfigManager', spec=AiShellConfigManager)
+    def test_success_when_config_file_available_with_chatgpt_config(
+        self,
+        mocked_config_manager_class: MagicMock,
+    ):
+        '''config file 이 있고, chatgpt_config 가 있는 경우 테스트 성공해야 한다.'''
+        # given
+        mocked_config_manager_class.is_config_file_available.return_value = True
+        mocked_config_manager_instance = mocked_config_manager_class.return_value
+        mocked_config_manager_instance.config_model =\
+            AiShellConfigModel(chatgpt_config=RevChatGPTChatbotConfigModel(session_token='invalid_session_token'))
+
+        # when
+        save_config('valid_session_token')
+
+        # then
+        mocked_config_manager_class.return_value.save_config.assert_called_once()
+
+    @patch('aishell.cli.config_aishell.AiShellConfigManager', spec=AiShellConfigManager)
+    def test_success_when_config_file_available_without_chatgpt_config(
+        self,
+        mocked_config_manager_class: MagicMock,
+    ):
+        '''config file 이 있고, chatgpt_config 가 없는 경우 테스트 성공해야 한다.'''
+        # given
+        mocked_config_manager_class.is_config_file_available.return_value = True
+        mocked_config_manager_instance = mocked_config_manager_class.return_value
+        mocked_config_manager_instance.config_model =\
+            AiShellConfigModel(language_model=LanguageModel.GPT3, openai_api_key='valid_openai_api_key')
+
+        # when
+        save_config('valid_session_token')
+
+        # then
+        mocked_config_manager_class.return_value.save_config.assert_called_once()
+
+    # when config file available - with chatgpt_config
+    # when config file available - without chatgpt_config
+    # when config file not available

aishell/models/revchatgpt_chatbot_config_model.py

@@ -18,6 +18,3 @@ def check_at_least_one_account_info(cls, values: dict[str, Optional[str]]):
             raise ValueError('No information for authentication provided.')
 
         return values
-
-    class Config:
-        frozen = True

aishell/query_clients/gpt3_client.py

@@ -11,6 +11,12 @@
 
 class GPT3Client(QueryClient):
 
+    def __init__(
+        self,
+        openai_api_key: str,
+    ):
+        openai.api_key = openai_api_key
+
     def query(self, prompt: str) -> str:
         prompt = self._construct_prompt(prompt)
         completion: Final[OpenAIResponseModel] = cast(
@@ -29,8 +35,9 @@ def query(self, prompt: str) -> str:
         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.
+        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: '''
+You: '''[1:]

aishell/query_clients/official_chatgpt_client.py

@@ -13,9 +13,8 @@ class OfficialChatGPTClient(QueryClient):
 
     def __init__(
         self,
-        openai_api_key: Optional[str] = None,
+        openai_api_key: str,
     ):
-        super().__init__()
         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')
@@ -31,6 +30,7 @@ def query(self, prompt: str) -> str:
         return executable_command
 
     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.'''
+        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.'''[1:]

aishell/query_clients/reverse_engineered_chatgpt_client.py

@@ -49,6 +49,7 @@ def query(self, prompt: str) -> str:
         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.'''
+        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.'''[1:]