Merge pull request #4 from code-yeongyu/feature/split

Author: code-yeongyu

aishell/models/__init__.py

@@ -0,0 +1 @@
+from .open_ai_response_model import OpenAIResponseModel as OpenAIResponseModel

aishell/models/open_ai_response_model.py

@@ -0,0 +1,24 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class OpenAIResponseModel(BaseModel):
+
+    class Choice(BaseModel):
+        finish_reason: str
+        index: int
+        logprobs: Optional[None]
+        text: Optional[str]
+
+    class Usage(BaseModel):
+        completion_tokens: int
+        prompt_tokens: int
+        total_tokens: int
+
+    choices: Optional[List[Choice]]
+    created: int
+    id: str
+    model: str
+    object: str
+    usage: Usage

aishell/py.typed

@@ -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/__init__.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/chatgpt_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/gpt3_client.py

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

aishell/query_clients/query_client.py

@@ -0,0 +1 @@
+from .make_executable_command import make_executable_command as make_executable_command

aishell/utils/__init__.py

@@ -0,0 +1,14 @@
+def make_executable_command(command: str) -> str:
+    # starting '\n' or trailing '\n' should be replaced as ''
+    # starting ' ' or trailing ' ' should be replaced as ''
+    if command.startswith('\n'):
+        command = command[1:]
+    if command.endswith('\n'):
+        command = command[:-1]
+    if command.startswith('`'):
+        command = command[1:]
+    if command.endswith('`'):
+        command = command[:-1]
+    command = command.strip()
+    command = command.split('User: ')[0]
+    return command

aishell/utils/make_executable_command.py

@@ -1,5 +1,5 @@
 [build-system]
-requires = [ "poetry-core",]
+requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
@@ -39,8 +39,8 @@ typeCheckingMode = "strict"
 stubPath = "./.type_stubs"
 venvPath = "./.venv"
 pythonPlatform = "All"
-include = [ "./*",]
-exclude = [ "**/node_modules", "**/__pycache__",]
+include = ["./*"]
+exclude = ["**/node_modules", "**/__pycache__"]
 useLibraryCodeForTypes = true
 verboseOutput = false
 reportMissingTypeStubs = false
@@ -55,7 +55,7 @@ allow_split_before_dict_value = false
 
 [tool.ruff]
 line-length = 119
-select = [ "PLC", "PLE", "PLR", "PLW", "E", "W", "F", "I", "Q", "C", "B",]
+select = ["PLE", "PLR", "PLW", "E", "W", "F", "I", "Q", "C", "B"]
 
 [tool.autoflake]
 check = false
@@ -71,7 +71,7 @@ remove-unused-variables = true
 python = "^3.8"
 poetry = "^1.3.1"
 revchatgpt = "^1.0.7"
-typer = {extras = ["all"], version = "^0.7.0"}
+typer = { extras = ["all"], version = "^0.7.0" }
 openai = "^0.26.5"
 pydantic = "^1.10.4"