From 1ebc39410be6c43ac4d91a5faa3eb0c3376c9f6c Mon Sep 17 00:00:00 2001
From: YeonGyu-Kim <code.yeon.gyu@gmail.com>
Date: Sun, 5 Mar 2023 22:03:08 +0900
Subject: [PATCH] =?UTF-8?q?Hello=20YGKA=20=F0=9F=99=8C?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ygka/cli/config_ygka.py | 71 +++++++++++++++++++++++++++++++++++++++++
 ygka/cli/ygka.py        | 55 +++++++++++++++++++++++++++++--
 2 files changed, 123 insertions(+), 3 deletions(-)
 create mode 100644 ygka/cli/config_ygka.py

diff --git a/ygka/cli/config_ygka.py b/ygka/cli/config_ygka.py
new file mode 100644
index 0000000..be003f4
--- /dev/null
+++ b/ygka/cli/config_ygka.py
@@ -0,0 +1,71 @@
+import sys
+
+import rich
+import typer
+from yt_dlp.cookies import SUPPORTED_BROWSERS
+
+from ygka.adapters.openai_cookie_adapter import OpenAICookieAdapter
+from ygka.models import RevChatGPTChatbotConfigModel
+from ygka.models.ygka_config_model import YGKAConfigModel
+from ygka.utils import YGKAConfigManager
+
+
+def config_ygka():
+    rich.print('''
+Hi! 🙌 I am [bold blue]YGKA[/bold blue]!
+[yellow][blue bold underline]Y[/blue bold underline]our
+[blue bold underline]G[/blue bold underline]enius,
+[blue bold underline]K[/blue bold underline]nowledgeable assistant
+[blue bold underline]A[/blue bold underline]n advanced ChatGPT client[/yellow] 🔥
+I am here to assist you with configuring YGKA. 💪
+
+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]YGKA[/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 = YGKAConfigManager.is_config_file_available(YGKAConfigManager.DEFAULT_CONFIG_PATH)
+    if is_config_file_available:
+        config_manager = YGKAConfigManager(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)
+        YGKA_config = YGKAConfigModel(chatgpt_config=chatgpt_config)
+        config_manager = YGKAConfigManager(config_model=YGKA_config)
+
+    config_manager.save_config()
+    return config_manager
diff --git a/ygka/cli/ygka.py b/ygka/cli/ygka.py
index b679d75..d8634d6 100644
--- a/ygka/cli/ygka.py
+++ b/ygka/cli/ygka.py
@@ -1,10 +1,59 @@
-from rich import print as pprint
+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 .cli_app import cli_app
+from .config_ygka import config_ygka
 from .retrieve_stdin import retrieve_stdin
 
 
 @cli_app.command()
-def ygka_command(query: str):
+def ygka_command(query: str, language_model: Optional[LanguageModel] = None):
     stdin = retrieve_stdin()
-    pprint({'stdin': stdin, 'query': query})
+
+    config_manager = _get_config_manager()
+    config_manager.config_model.language_model = language_model or config_manager.config_model.language_model
+
+    query_client = QueryClientFactory(config_model=config_manager.config_model).create()
+
+    console = Console()
+    try:
+        with console.status('''[green] YGKA is waiting for ChatGPT's answer ...[/green]'''):
+            prompt = _generate_prompt(stdin, query)
+            response = query_client.query(prompt)
+            console.print(response)
+    except KeyError:
+        rich.print('It looks like the `session_token` is expired. Please reconfigure YGKA.')
+        typer.confirm('Reconfigure YGKA?', abort=True)
+        config_ygka()
+        ygka_command(query=query, language_model=language_model)
+        typer.Exit()
+
+
+def _get_config_manager():
+    is_config_file_available = YGKAConfigManager.is_config_file_available(YGKAConfigManager.DEFAULT_CONFIG_PATH)
+    if is_config_file_available:
+        return YGKAConfigManager(load_config=True)
+    else:
+        return config_ygka()
+
+
+def _generate_prompt(stdin: Optional[str], prompt: str) -> str:
+    if stdin:
+        return f'''
+stdin:
+{stdin}
+
+prompt:
+{prompt}
+
+[system] don't mention that i gave you in a format of stdin and prompt
+'''[1:]
+    else:
+        return prompt