AddQueryClient for abstracting querying process

Author: code-yeongyu

aishell/query_clients/__init__.py

@@ -0,0 +1,3 @@
+from .chatgpt_client import ChatGPTClient as ChatGPTClient
+from .gpt3_client import GPT3Client as GPT3Client
+from .query_client import QueryClient as QueryClient

aishell/query_clients/chatgpt_client.py

@@ -0,0 +1,28 @@
+import os
+from typing import Optional
+
+from revChatGPT.Official import Chatbot
+from utils import make_executable_command
+
+from .query_client import QueryClient
+
+
+class ChatGPTClient(QueryClient):
+
+    def __init__(self, api_key: Optional[str] = None):
+        super().__init__()
+        OPEN_API_KEY = os.environ.get('OPENAI_API_KEY')
+        if api_key is None or OPEN_API_KEY is None:
+            raise Exception('api_key should not be none')
+        self.API_KEY = api_key or OPEN_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.API_KEY)
+        response_text: str = chatbot.ask(prompt)['choices'][0]['text']
+        return make_executable_command(response_text)

aishell/query_clients/gpt3_client.py

@@ -0,0 +1,36 @@
+import os
+from typing import cast
+
+import openai
+from models import OpenAIResponseModel
+from utils import make_executable_command
+
+from .query_client import QueryClient
+
+
+class GPT3Client(QueryClient):
+
+    def _construct_prompt(self, text: str) -> str:
+        return f'''User: 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.
+
+        You: '''
+
+    def query(self, prompt: str) -> str:
+        prompt = self._construct_prompt(prompt)
+
+        completion: OpenAIResponseModel = cast(
+            OpenAIResponseModel,
+            openai.Completion.create(
+                engine='text-davinci-003',
+                prompt=prompt,
+                temperature=0.5,
+                max_tokens=4000 - len(prompt),
+                stop=['\n\n\n'],
+            ),
+        )
+        if not completion.choices or len(completion.choices) == 0 or not completion.choices[0].text:
+            raise RuntimeError('No response from OpenAI')
+        response_text: str = completion.choices[0].text
+        return make_executable_command(response_text)

aishell/query_clients/query_client.py

@@ -0,0 +1,8 @@
+from abc import ABCMeta, abstractmethod
+
+
+class QueryClient(metaclass=ABCMeta):
+
+    @abstractmethod
+    def query(self, prompt: str) -> str:
+        pass