Merge pull request #26 from code-yeongyu/feature/cookie-parsing

Author: code-yeongyu

.github/workflows/test.yml

@@ -0,0 +1,34 @@
+name: Test
+on:
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+      - name: Set up Python 3.9
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.9
+
+      - uses: snok/[email protected]
+        with:
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+
+      - name: Cache Dependencies
+        uses: actions/cache@v2
+        id: cache-dependencies
+        with:
+          path: .venv
+          key: venv-${{ runner.os }}-${{ hashFiles('./poetry.lock') }}
+
+      - name: Install Dependencies if cache doesn't hit
+        if: steps.cache-dependencies.cache-hit != 'true'
+        run: poetry install
+
+      - name: Run test
+        run: poetry run invoke test

.vscode/launch.json

@@ -0,0 +1,18 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "AiShell Debug",
+            "type": "python",
+            "request": "launch",
+            "module": "aishell",
+            "justMyCode": true,
+            "args": [
+                "print helloworld"
+            ]
+        }
+    ]
+}

README.md

@@ -38,10 +38,15 @@ aishell --help
 
 ### For those who want to use reverse-engineered `ChatGPT`
 
-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_KEY`
-1. Enjoy AiShell
+- 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
 
 ### For those who want to use `Official ChatGPT(GPT3.5-turbo)` or `GPT-3`
 

aishell/adapters/__init__.py

@@ -0,0 +1 @@
+from .openai_chatgpt_adapter import OpenAIChatGPTAdapter as OpenAIChatGPTAdapter

aishell/adapters/openai_chatgpt_adapter.py

@@ -0,0 +1,27 @@
+from http.cookiejar import Cookie
+from typing import Optional
+
+from yt_dlp.cookies import SUPPORTED_BROWSERS, extract_cookies_from_browser
+from yt_dlp.utils import YoutubeDLCookieJar
+
+
+class OpenAIChatGPTAdapter:
+
+    def __init__(self, browser_name: str):
+        if browser_name not in SUPPORTED_BROWSERS:  # type: ignore
+            raise ValueError(f'Browser {browser_name} is not supported. Supported browsers are: {SUPPORTED_BROWSERS}')
+        self.BROWSER_NAME = browser_name
+
+    def get_openai_cookies(self) -> dict[str, Optional[str]]:
+        all_cookies: YoutubeDLCookieJar = extract_cookies_from_browser(self.BROWSER_NAME)
+        try:
+            openai_cookies: dict[str, Cookie] = all_cookies.__dict__['_cookies']['chat.openai.com']['/']
+        except KeyError as error:
+            raise ValueError('Could not find OpenAI cookies. Make sure you are logged in to OpenAI.') from error
+
+        cookies: dict[str, Optional[str]] = {key: cookie.value for key, cookie in openai_cookies.items()}
+
+        return cookies
+
+    def get_openai_session_token(self) -> Optional[str]:
+        return self.get_openai_cookies()['__Secure-next-auth.session-token']

aishell/adapters/test/__init__.py

@@ -0,0 +1,17 @@
+import pytest
+from yt_dlp.cookies import SUPPORTED_BROWSERS
+
+from aishell.adapters import OpenAIChatGPTAdapter
+
+
+class TestChatGPTAccessTokenAdapterInit:
+
+    def test_fail_when_invalid_browser_name_is_given(self) -> None:
+        '''잘못된 브라우저 이름을 입력했을 때 ValueError가 발생한다.'''
+        with pytest.raises(ValueError):
+            OpenAIChatGPTAdapter('invalid_browser_name')
+
+    @pytest.mark.parametrize('browser_name', SUPPORTED_BROWSERS)
+    def test_success_when_valid_browser_name_is_given(self, browser_name: str) -> None:
+        '''올바른 브라우저 이름을 입력했을 때 ValueError가 발생하지 않는다'''
+        OpenAIChatGPTAdapter(browser_name)

aishell/adapters/test/test_openai_chatgpt_adapter.py

@@ -1,23 +1,40 @@
 import os
+import webbrowser
 
 import typer
 from rich.console import Console
+from yt_dlp.cookies import SUPPORTED_BROWSERS
 
+from aishell.adapters.openai_chatgpt_adapter import OpenAIChatGPTAdapter
+from aishell.exceptions import UnauthorizedAccessError
 from aishell.models.language_model import LanguageModel
 from aishell.query_clients import GPT3Client, OfficialChatGPTClient, QueryClient, ReverseEngineeredChatGPTClient
 
 cli_app = typer.Typer()
 
 
+def _open_chatgpt_browser():
+    CHATGPT_LOGIN_URL = 'https://chat.openai.com/auth/login?next=/chat'
+    webbrowser.open(CHATGPT_LOGIN_URL)
+
+
 @cli_app.command()
-def ask(question: str, language_model: LanguageModel = LanguageModel.OFFICIAL_CHATGPT):
+def ask(question: str, language_model: LanguageModel = LanguageModel.REVERSE_ENGINEERED_CHATGPT):
     query_client: QueryClient
     if language_model == LanguageModel.GPT3:
         query_client = GPT3Client()
     elif language_model == LanguageModel.OFFICIAL_CHATGPT:
         query_client = OfficialChatGPTClient()
     elif language_model == LanguageModel.REVERSE_ENGINEERED_CHATGPT:
-        query_client = ReverseEngineeredChatGPTClient()
+        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 = OpenAIChatGPTAdapter(BROWSER_NAME)
+            session_token = adapter.get_openai_session_token()
+            query_client = ReverseEngineeredChatGPTClient(session_token=session_token)
 
     query_client.query(question)
 

aishell/cli.py

@@ -0,0 +1 @@
+from .unauthorized_access_error import UnauthorizedAccessError as UnauthorizedAccessError

aishell/exceptions/__init__.py

@@ -0,0 +1,2 @@
+class UnauthorizedAccessError(Exception):
+    pass