Add OfficialChatGPTClient

Author: code-yeongyu

aishell/query_clients/__init__.py

@@ -1,3 +1,4 @@
 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)