Merge pull request #29 from code-yeongyu/feature/improve-readibility

Author: code-yeongyu

aishell/models/__init__.py

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

aishell/models/open_ai_response_model.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import Optional
 
 from pydantic import BaseModel
 
@@ -16,9 +16,12 @@ class Usage(BaseModel):
         prompt_tokens: int
         total_tokens: int
 
-    choices: Optional[List[Choice]]
+    choices: Optional[list[Choice]]
     created: int
     id: str
     model: str
     object: str
     usage: Usage
+
+    class Config:
+        frozen = True

aishell/models/revchatgpt_chatbot_config_model.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+from pydantic import BaseModel, root_validator
+
+
+class RevChatGPTChatbotConfigModel(BaseModel):
+    email: Optional[str] = None
+    password: Optional[str] = None
+    session_token: Optional[str] = None
+    access_token: Optional[str] = None
+    paid: bool = False
+
+    @root_validator
+    def check_at_least_one_account_info(cls, values: dict[str, Optional[str]]):
+        IS_ACCOUNT_LOGIN = values.get('email') and values.get('password')
+        IS_TOKEN_AUTH = values.get('session_token') or values.get('access_token')
+        if not IS_ACCOUNT_LOGIN and not IS_TOKEN_AUTH:
+            raise ValueError('No information for authentication provided.')
+
+        return values
+
+    class Config:
+        frozen = True

aishell/query_clients/gpt3_client.py

@@ -1,5 +1,5 @@
 import os
-from typing import cast
+from typing import Final, cast
 
 import openai
 
@@ -11,16 +11,9 @@
 
 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(  # type: ignore [no-any-unimported]
+        completion: Final[OpenAIResponseModel] = cast(
             OpenAIResponseModel,
             openai.Completion.create(
                 engine='text-davinci-003',
@@ -32,5 +25,12 @@ def query(self, prompt: str) -> str:
         )
         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
+        response_text: Final[str] = completion.choices[0].text
         return make_executable_command(response_text)
+
+    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: '''

aishell/query_clients/official_chatgpt_client.py

@@ -1,5 +1,5 @@
 import os
-from typing import Final, Optional
+from typing import Optional
 
 from revChatGPT.V3 import Chatbot
 
@@ -10,28 +10,27 @@
 
 
 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)
+        OPENAI_API_KEY: Optional[str] = os.environ.get('OPENAI_API_KEY', openai_api_key)
         if OPENAI_API_KEY is None:
             raise UnauthorizedAccessError('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.'''
+        self.OPENAI_API_KEY = OPENAI_API_KEY
 
     def query(self, prompt: str) -> str:
-        prompt = self._construct_prompt(prompt)
+        chatbot = Chatbot(api_key=self.OPENAI_API_KEY)
 
-        chatbot = Chatbot(api_key=self.openai_api_key)
+        prompt = self._construct_prompt(prompt)
         response_text = chatbot.ask(prompt)
+        executable_command = make_executable_command(response_text)
+        return executable_command
 
-        return make_executable_command(response_text)
+    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.'''

aishell/query_clients/reverse_engineered_chatgpt_client.py

@@ -1,16 +1,21 @@
 import os
-from typing import Optional, cast
+from typing import Optional, Union, cast
 
 from revChatGPT.V1 import Chatbot
 
 from aishell.exceptions import UnauthorizedAccessError
+from aishell.models import RevChatGPTChatbotConfigModel
 from aishell.utils import make_executable_command
 
 from .query_client import QueryClient
 
 
 class ReverseEngineeredChatGPTClient(QueryClient):
-    config: dict[str, str] = {}
+    _config: RevChatGPTChatbotConfigModel
+
+    @property
+    def revchatgpt_config(self) -> dict[str, Union[str, bool]]:
+        return self._config.dict(exclude_none=True)
 
     def __init__(
         self,
@@ -19,21 +24,17 @@ def __init__(
     ):
         CHATGPT_ACCESS_TOKEN = os.environ.get('CHATGPT_ACCESS_TOKEN', access_token)
         CHATGPT_SESSION_TOKEN = os.environ.get('CHATGPT_SESSION_TOKEN', session_token)
-        if CHATGPT_ACCESS_TOKEN is not None:
-            self.config['access_token'] = CHATGPT_ACCESS_TOKEN
-        elif CHATGPT_SESSION_TOKEN is not None:
-            self.config['session_token'] = CHATGPT_SESSION_TOKEN
+        if CHATGPT_ACCESS_TOKEN:
+            self._config = RevChatGPTChatbotConfigModel(access_token=CHATGPT_ACCESS_TOKEN)
+        elif CHATGPT_SESSION_TOKEN:
+            self._config = RevChatGPTChatbotConfigModel(session_token=CHATGPT_SESSION_TOKEN)
         else:
             raise UnauthorizedAccessError('No access token or session token provided.')
 
-    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(config=self.config)
+        chatbot = Chatbot(config=self.revchatgpt_config)  #  pyright: ignore [reportGeneralTypeIssues]
+        # ignore for wrong type hint of revchatgpt
 
         response_text = ''
         for data in chatbot.ask(prompt):
@@ -42,3 +43,8 @@ def query(self, prompt: str) -> str:
         response_text = make_executable_command(cast(str, response_text))
 
         return response_text
+
+    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.'''

aishell/utils/str_enum.py

@@ -4,7 +4,12 @@
 
 class StrEnum(str, Enum):
 
-    def _generate_next_value_(name: str, start: int, count: int, last_values: list[Any]):  # type: ignore
+    def _generate_next_value_(  # pyright: ignore [reportIncompatibleMethodOverride], for pyright's bug
+        name: str,
+        start: int,
+        count: int,
+        last_values: list[Any],
+    ):
         return name.lower()
 
     def __repr__(self):