Merge pull request #25 from code-yeongyu/feature/chatgpt-official

Add support for Official ChatGPT API 🚀

Author: code-yeongyu

README.md

@@ -43,17 +43,13 @@ aishell --help
 1. Set the API key as an environment variable `CHATGPT_ACCESS_KEY`
 1. Enjoy AiShell
 
-### For those who want to use `GPT-3`
+### 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
 
-### For those who want to use Official ChatGPT API `gpt-3.5-turbo`
-
-- Currently not supported, but soon will be supported!
-
 ## Contributions 💬
 
 Feel free to contribute to AiShell by adding more functionality or fixing bugs.

aishell/cli.py

@@ -3,18 +3,23 @@
 import typer
 from rich.console import Console
 
-from aishell.query_clients import ChatGPTClient, GPT3Client, QueryClient
+from aishell.models.language_model import LanguageModel
+from aishell.query_clients import GPT3Client, OfficialChatGPTClient, QueryClient, ReverseEngineeredChatGPTClient
 
 cli_app = typer.Typer()
 
 
 @cli_app.command()
-def ask(question: str, use_chatgpt: bool = False):
+def ask(question: str, language_model: LanguageModel = LanguageModel.OFFICIAL_CHATGPT):
     query_client: QueryClient
-    if use_chatgpt:
-        query_client = ChatGPTClient()
-    else:
+    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()
+
+    query_client.query(question)
 
     console = Console()
     with console.status(

aishell/models/__init__.py

@@ -1 +1,2 @@
+from .language_model import LanguageModel as LanguageModel
 from .open_ai_response_model import OpenAIResponseModel as OpenAIResponseModel

aishell/models/language_model.py

@@ -0,0 +1,9 @@
+from enum import auto
+
+from aishell.utils import StrEnum
+
+
+class LanguageModel(StrEnum):
+    OFFICIAL_CHATGPT = auto()
+    REVERSE_ENGINEERED_CHATGPT = auto()
+    GPT3 = auto()

aishell/query_clients/__init__.py

@@ -1,3 +1,4 @@
-from .chatgpt_client import ChatGPTClient as ChatGPTClient
 from .gpt3_client import GPT3Client as GPT3Client
+from .official_chatgpt_client import OfficialChatGPTClient as OfficialChatGPTClient
 from .query_client import QueryClient as QueryClient
+from .reverse_engineered_chatgpt_client import ReverseEngineeredChatGPTClient as ReverseEngineeredChatGPTClient

aishell/query_clients/official_chatgpt_client.py

@@ -0,0 +1,36 @@
+import os
+from typing import Final, Optional
+
+from revChatGPT.V3 import Chatbot
+
+from aishell.utils import make_executable_command
+
+from .query_client import QueryClient
+
+
+class OfficialChatGPTClient(QueryClient):
+    openai_api_key: str
+
+    def __init__(
+        self,
+        openai_api_key: Optional[str] = None,
+    ):
+        super().__init__()
+        OPENAI_API_KEY: Final[Optional[str]] = os.environ.get('OPENAI_API_KEY', openai_api_key)
+        if OPENAI_API_KEY is None:
+            raise Exception('OPENAI_API_KEY should not be none')
+
+        self.openai_api_key = OPENAI_API_KEY
+
+    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.'''
+
+    def query(self, prompt: str) -> str:
+        prompt = self._construct_prompt(prompt)
+
+        chatbot = Chatbot(api_key=self.openai_api_key)
+        response_text = chatbot.ask(prompt)
+
+        return make_executable_command(response_text)

aishell/query_clients/chatgpt_client.py renamed to aishell/query_clients/reverse_engineered_chatgpt_client.py

@@ -1,5 +1,5 @@
 import os
-from typing import Optional
+from typing import Optional, cast
 
 from revChatGPT.V1 import Chatbot
 
@@ -8,7 +8,7 @@
 from .query_client import QueryClient
 
 
-class ChatGPTClient(QueryClient):
+class ReverseEngineeredChatGPTClient(QueryClient):
     access_key: str
 
     def __init__(
@@ -38,6 +38,6 @@ def query(self, prompt: str) -> str:
         for data in chatbot.ask(prompt):
             response_text = data['message']
 
-        response_text = make_executable_command(response_text)
+        response_text = make_executable_command(cast(str, response_text))
 
         return response_text

aishell/utils/__init__.py

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

aishell/utils/str_enum.py

@@ -0,0 +1,14 @@
+from enum import Enum
+from typing import Any
+
+
+class StrEnum(str, Enum):
+
+    def _generate_next_value_(name: str, start: int, count: int, last_values: list[Any]):  # type: ignore
+        return name.lower()
+
+    def __repr__(self):
+        return self.name
+
+    def __str__(self):
+        return self.name