Merge pull request #2 from RexWzh/add_config

Use config files

.env.example

@@ -0,0 +1,14 @@
+#!/bin/bash
+# Description: Env file for askchat.
+# Current version: 0.1.1
+
+# The base url of the API (without suffix /v1)
+OPENAI_API_BASE_URL='localhost:8000'
+
+# Your API key
+OPENAI_API_KEY='EMPTY'
+
+# The model name.
+# You can use `askchat --all-valid-models` to see the valid models
+OPENAI_API_MODEL='baichuan2'
+

README.md

@@ -23,7 +23,12 @@ ask hello
 Ask with more options via `askchat`:
 ```bash
 # ask with a specific model
-askchat hello -m "gpt-4"
+askchat hello -m "baichuan2" --base_url "localhost:8000"
+```
+
+Generate config file for default options:
+```bash
+askchat --generate-config
 ```
 
 Other options:
@@ -32,6 +37,8 @@ Other options:
 askchat -v 
 # Get debug log
 askchat --debug
-# get valid models
+# get valid models that contains "gpt"
 askchat --valid-models
+# get all valid models
+askchat --all-valid-models
 ```

askchat/__init__.py

@@ -2,6 +2,6 @@
 
 __author__ = """Rex Wang"""
 __email__ = '[email protected]'
-__version__ = '0.1.1'
+__version__ = '0.2.0'
 
 from .askchat import ask

askchat/askchat.py

@@ -1,16 +1,26 @@
 """Main module."""
-from chattool import Chat, debug_log
-import asyncio, os, uuid
+
 from argparse import ArgumentParser
-import askchat
 from pprint import pprint
+from dotenv import load_dotenv, set_key
+import asyncio, os, uuid
+import askchat
+
 VERSION = askchat.__version__
-CONFIG_FILE = os.path.expanduser("~/.askrc")
+CONFIG_PATH = os.path.expanduser("~/.askchat")
+CONFIG_FILE = os.path.expanduser("~/.askchat/.env")
+## read para from config file
+if os.path.exists(CONFIG_FILE):
+    load_dotenv(CONFIG_FILE, override=True)
+
+# load chattool after update the config
+from chattool import Chat, debug_log
 
 # print the response in a typewriter way
-async def show_resp(chat):
+async def show_resp(chat, delay=0.01):
     async for char in chat.async_stream_responses(textonly=True):
         print(char, end='', flush=True)
+        await asyncio.sleep(delay)
 
 def ask():
     """Interact with ChatGPT in terminal via chattool"""
@@ -30,21 +40,24 @@ def main():
     """Interact with ChatGPT in terminal via chattool"""
     # parse arguments
     parser = ArgumentParser()
-    ## use nargs='?' to make message optional
+    ## arguments for chat message
     parser.add_argument('message', help='User message', default='', nargs='*')
-    parser.add_argument('-v', '--version', action='version', version=VERSION)
-    parser.add_argument('--debug', action='store_true', help='Print debug log')
-    parser.add_argument('--valid-models', action='store_true', help='Print valid models that contain "gpt" in their names')
-    parser.add_argument('--all-valid-models', action='store_true', help='Print all valid models')
     parser.add_argument('-m', '--model', default=None, help='Model name')
     parser.add_argument('--base-url', default=None, help='base url of the api(without suffix `/v1`)')
     parser.add_argument("--api-key", default=None, help="API key")
-    parser.add_argument('--generate-config', action="store_true", help="Generate a configuration file by environment table.")
+    ## other options
+    parser.add_argument('--debug', action='store_true', help='Print debug log')
+    parser.add_argument('--valid-models', action='store_true', help='Print valid models that contain "gpt" in their names')
+    parser.add_argument('--all-valid-models', action='store_true', help='Print all valid models')
+    parser.add_argument('--generate-config', action="store_true", help="Generate a configuration file by environment table")
+    parser.add_argument('-v', '--version', action='version', version=VERSION)
     args = parser.parse_args()
+
     # show debug log
     if args.debug:
         debug_log()
         return
+
     # show valid models
     if args.valid_models:
         print('Valid models that contain "gpt" in their names:')
@@ -54,52 +67,47 @@ def main():
         print('All valid models:')
         pprint(Chat().get_valid_models(gpt_only=False))
         return
-    if args.generate_config is not None:
-        api_key = os.environ.get("OPENAI_API_KEY", "")
-        base_url = os.environ.get("OPENAI_API_BASE_URL", "")
-        model = os.environ.get("OPENAI_API_MODEL", "")
+
+    # generate config file
+    if args.generate_config:
+        api_key = os.environ.get("OPENAI_API_KEY")
+        base_url = os.environ.get("OPENAI_API_BASE_URL")
+        model = os.environ.get("OPENAI_API_MODEL")
+        # move the old config file to a temporary file
         if os.path.exists(CONFIG_FILE):
             # create a temporary file
             os.makedirs("/tmp", exist_ok=True)
-            tmp_file = os.path.join("/tmp", str(uuid.uuid4())[:8] + ".askrc")
+            tmp_file = os.path.join("/tmp", str(uuid.uuid4())[:8] + ".askchat.env")
             # move the old config file to a temporary file
             os.rename(CONFIG_FILE, tmp_file)
             print(f"Moved old config file to {tmp_file}")
+        # save the config file
+        os.makedirs(CONFIG_PATH, exist_ok=True)
         with open(CONFIG_FILE, "w") as f:
             # description for the config file
             f.write("#!/bin/bash\n" +\
-                    "# Description: This is a configuration file for askchat.\n" +\
-                    "# Author: Rex Wang\n" +\
-                    "# Current version: " + VERSION + "\n\n")
-            # write the environment table
-            f.write("# Your API key\n")
-            f.write(f"OPENAI_API_KEY={api_key}\n\n")
-            f.write("# The base url of the API (without suffix /v1)\n")
-            f.write(f"OPENAI_API_BASE_URL={base_url}\n\n")
-            f.write("# The model name. You can use `askchat --all-valid-models` to see the valid models.\n")
-            f.write(f"OPENAI_API_MODEL={model}\n\n")
-            print("Created config file at", CONFIG_FILE)
+                    "# Description: Env file for askchat.\n" +\
+                    "# Current version: " + VERSION + "\n\n" +\
+                    "# The base url of the API (without suffix /v1)\n" +\
+                    "OPENAI_API_BASE_URL=\n\n" +\
+                    "# Your API key\n" +\
+                    "OPENAI_API_KEY=\n\n" +\
+                    "# The model name\n" +\
+                    "# You can use `askchat --all-valid-models` to see the valid models\n" +\
+                    "OPENAI_API_MODEL=\n\n")
+        # write the environment table
+        if api_key: set_key(CONFIG_FILE, "OPENAI_API_KEY", api_key)
+        if base_url: set_key(CONFIG_FILE, "OPENAI_API_BASE_URL", base_url)
+        if model: set_key(CONFIG_FILE, "OPENAI_API_MODEL", model)
+        print("Created config file at", CONFIG_FILE)
         return
 
     # get message, model, and base url
     msg = args.message
     if isinstance(msg, list):
         msg = ' '.join(msg)
     assert len(msg.strip()), 'Please specify message'
-    # read para from config or args
-    if os.path.exists(CONFIG_FILE):
-        with open(CONFIG_FILE) as f:
-            lines = f.readlines()
-        for line in lines:
-            if line.startswith("OPENAI_API_KEY="):
-                api_key = line.split("=")[-1].strip()
-            elif line.startswith("OPENAI_API_BASE_URL="):
-                base_url = line.split("=")[-1].strip()
-            elif line.startswith("OPENAI_API_MODEL="):
-                model = line.split("=")[-1].strip()
-    api_key = args.api_key if hasattr(args, "api_key") else api_key
-    base_url = args.base_url if hasattr(args, "base_url") else base_url
-    model = args.model if hasattr(args, "model") else model
+
     # call the function
-    chat = Chat(msg, model=model, base_url=base_url, api_key=api_key)
+    chat = Chat(msg, model=args.model, base_url=args.base_url, api_key=args.api_key)
     asyncio.run(show_resp(chat))

setup.py

@@ -4,12 +4,12 @@
 
 from setuptools import setup, find_packages
 
-VERSION = '0.1.1'
+VERSION = '0.2.0'
 
 with open('README.md') as readme_file:
     readme = readme_file.read()
 
-requirements = ['chattool>=2.5.0']
+requirements = ['chattool>=2.5.0', "python-dotenv>=0.17.0"]
 
 test_requirements = ['pytest>=3']
 
@@ -26,6 +26,7 @@
         'Programming Language :: Python :: 3.6',
         'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
     ],
     description="Interact with ChatGPT in terminal via chattool",
     install_requires=requirements,