Merge pull request #31 from code-yeongyu/feature/configs

Author: code-yeongyu

README.md

@@ -16,44 +16,42 @@ A simple Python code that connects to OpenAI's ChatGPT and executes the returned
 - Automatically executes the command from the response of ChatGPT
 - Good for complex tasks like handling Git and extracting tar files
 - No need to search StackOverflow for commands, AiShell has got you covered
+- **No need to set up annoying retrieving of tokens or API keys with ChatGPT, as AiShell does it for you. INSTALL IT. EXECUTE IT. DONE.**
+
+## Prerequisites 📚
+
+- Python 3.9+
+- ChatGPT Account (or OpenAI Account)
 
 ## Installation 🔧
 
 ```sh
 pip install aishell
 ```
 
-## Usage 📝
+## Getting Started 🚀
+
+Let's just start by printing "Hello World" using AiShell.
 
 ```sh
-aishell --help
+aishell 'print Hello World'
 ```
 
-## Prerequisites 📚
-
-- Python 3
-- OpenAI API Key or ChatGPT Account
-
-## Getting Started 🚀
-
-### For those who want to use reverse-engineered `ChatGPT`
-
-- Permanent Login Method
-    1. Login on <https://chat.openai.com/>
-    1. Get your 'accessToken` from <https://chat.openai.com/api/auth/session>
-    1. Set the API key as an environment variable `CHATGPT_ACCESS_TOKEN`
-- Temporary Login Method
-    1. Just run `aishell <query>`
-    1. Browser opens up. Login there.
-    1. Tell AiShell which browser you use.
-    1. Enjoy AiShell
+## Advanced Settings 🛠
 
 ### For those who want to use `Official ChatGPT(GPT3.5-turbo)` or `GPT-3`
 
 1. Create account on OpenAI
 1. Go to <https://platform.openai.com/account/api-keys>, Copy API key
-1. Set the API key as an environment variable `OPENAI_API_KEY`
-1. Enjoy AiShell
+1. Modify or create `~/.aishell/config.json` file like following
+
+    ```sh
+    {
+        ...
+        "language_model": <language model of your preference>, //"official_chatgpt" or "gpt3"
+        "openai_api_key": <your openai api key>
+    }
+    ```
 
 ## Contributions 💬
 

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/models/__init__.py

@@ -1,3 +1,4 @@
+from .aishell_config_model import AiShellConfigModel as AiShellConfigModel
 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/aishell_config_model.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+from pydantic import BaseModel, root_validator
+
+from .language_model import LanguageModel
+from .revchatgpt_chatbot_config_model import RevChatGPTChatbotConfigModel
+
+
+class AiShellConfigModel(BaseModel):
+    language_model: LanguageModel = LanguageModel.REVERSE_ENGINEERED_CHATGPT
+    chatgpt_config: Optional[RevChatGPTChatbotConfigModel] = None
+    openai_api_key: Optional[str] = None
+
+    @root_validator
+    def check_required_info_provided(cls, values: dict[str, Optional[str]]):
+        OPENAI_API_KEY_REQUIRED_MODELS = (LanguageModel.GPT3, LanguageModel.OFFICIAL_CHATGPT)
+
+        language_model = values.get('language_model')
+        if language_model in OPENAI_API_KEY_REQUIRED_MODELS:
+            if not values.get('openai_api_key'):
+                raise ValueError('openai_api_key should not be none')
+        elif language_model == LanguageModel.REVERSE_ENGINEERED_CHATGPT:
+            if not values.get('chatgpt_config'):
+                raise ValueError('chatgpt_config should not be none')
+
+        return values

aishell/models/language_model.py

@@ -1,6 +1,6 @@
 from enum import auto
 
-from aishell.utils import StrEnum
+from .str_enum import StrEnum
 
 
 class LanguageModel(StrEnum):

aishell/utils/str_enum.py renamed to aishell/models/str_enum.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)

aishell/query_clients/reverse_engineered_chatgpt_client.py

@@ -1,2 +1,2 @@
+from .aishell_config_manager import AiShellConfigManager as AiShellConfigManager
 from .make_executable_command import make_executable_command as make_executable_command
-from .str_enum import StrEnum as StrEnum

aishell/utils/__init__.py

@@ -0,0 +1,48 @@
+import json
+import os
+from typing import Optional
+
+from aishell.models import AiShellConfigModel
+
+
+class AiShellConfigManager:
+    DEFAULT_CONFIG_PATH = os.path.expanduser('~/.aishell_config.json')
+
+    def __init__(
+        self,
+        config_model: Optional[AiShellConfigModel] = None,
+        config_path: Optional[str] = None,
+        load_config: bool = False,
+    ):
+        '''
+        Initialize an instance of AiShellConfigManager
+
+        Args:
+            config_model: \
+                    An instance of AiShellConfigManager to use as the configuration.\
+                    If None and load_config is True, loads the configuration from the configuration file.
+            config_path: Path to the configuration file. \
+                If None, uses the default path "~/.aishell_config.json".
+            load_config: If True and config_model is None, loads the configuration from the configuration file.
+        '''
+        self.config_path = config_path or AiShellConfigManager.DEFAULT_CONFIG_PATH
+
+        if config_model:
+            self.update_config(config_model)
+        elif load_config:
+            self.load_config()
+
+    @staticmethod
+    def is_config_file_available(config_path: str) -> bool:
+        return os.path.isfile(config_path)
+
+    def update_config(self, config_model: AiShellConfigModel):
+        self.config_model = config_model
+
+    def load_config(self):
+        self.config_model = AiShellConfigModel.parse_file(self.config_path)
+
+    def save_config(self):
+        with open(self.config_path, 'w') as aishell_config_path:
+            config_dict = self.config_model.dict()
+            aishell_config_path.write(json.dumps(config_dict, indent=4, sort_keys=True))