Replace AiShell's common logic with ygka

Author: code-yeongyu

aishell/cli/aishell.py

@@ -2,14 +2,13 @@
 import time
 from typing import Optional
 
-import rich
 import typer
 from rich.console import Console
+from ygka.models import LanguageModel
+from ygka.query_clients import QueryClientFactory
+from ygka.utils import YGKAConfigManager
 
-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 aishell.utils import construct_prompt
 
 from .cli_app import cli_app
 from .config_aishell import config_aishell
@@ -20,53 +19,29 @@ def aishell_command(question: str, language_model: Optional[LanguageModel] = Non
     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,
-    )
+    query_client = QueryClientFactory(config_model=config_manager.config_model).create()
 
     console = Console()
 
-    try:
-        with console.status(f'''
+    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()
+        start_time = time.time()
+        response = query_client.query(construct_prompt(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]')
+    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()
+    will_execute = typer.confirm('Execute this command?')
+    if will_execute:
+        os.system(response)
 
 
 def _get_config_manager():
-    is_config_file_available = AiShellConfigManager.is_config_file_available(AiShellConfigManager.DEFAULT_CONFIG_PATH)
+    is_config_file_available = YGKAConfigManager.is_config_file_available(YGKAConfigManager.DEFAULT_CONFIG_PATH)
     if is_config_file_available:
-        return AiShellConfigManager(load_config=True)
+        return YGKAConfigManager(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/config_aishell.py

@@ -2,15 +2,12 @@
 
 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
+from revChatGPTAuth import SupportedBrowser
+from ygka.cli.config_ygka import save_config
 
 
 def config_aishell():
+    SUPPORTED_BROWSERS = [browser.value for browser in SupportedBrowser]
     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. 💪
@@ -29,13 +26,7 @@ def config_aishell():
         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)
+    config_manager = save_config(browser_name=browser_name)
 
     rich.print(f'''
 [green bold]Excellent![/green bold] You are now ready to use [bold blue]AiShell[/bold blue] 🚀
@@ -46,22 +37,3 @@ def config_aishell():
 
 '''[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/utils/__init__.py

@@ -1,2 +1,3 @@
 from .aishell_config_manager import AiShellConfigManager as AiShellConfigManager
+from .construct_prompt import construct_prompt as construct_prompt
 from .make_executable_command import make_executable_command as make_executable_command

aishell/utils/construct_prompt.py

@@ -0,0 +1,10 @@
+import os
+
+
+def construct_prompt(query: str) -> str:
+    return f'''
+Translate the following task to a {os.uname()[0]} shell command. Users provide a text-query as input.
+Provide ONLY the command in ONE LINE, with no explanation:
+
+One-line command for: {query}
+'''[1:]