From 0b7d7d1251e7f444391a381d5c8f2f8304e96b12 Mon Sep 17 00:00:00 2001
From: rex <1073853456@qq.com>
Date: Tue, 12 Mar 2024 17:14:56 +0800
Subject: [PATCH] add autocompletion for askenv

---
 askchat/__init__.py  | 21 ++++++++++++++++++++-
 askchat/askenv.py    | 14 ++++++++------
 askchat/cli.py       | 18 ++----------------
 setup.py             |  4 ++--
 tests/test_askenv.py |  1 -
 5 files changed, 32 insertions(+), 26 deletions(-)

diff --git a/askchat/__init__.py b/askchat/__init__.py
index 4c09e09..1b394e0 100644
--- a/askchat/__init__.py
+++ b/askchat/__init__.py
@@ -2,10 +2,11 @@
 
 __author__ = """Rex Wang"""
 __email__ = '1073853456@qq.com'
-__version__ = '1.1.0'
+__version__ = '1.1.1'
 
 import asyncio
 from pathlib import Path
+import click
 
 # Main environment file
 CONFIG_PATH = Path.home() / ".askchat"
@@ -13,6 +14,24 @@
 MAIN_ENV_PATH = Path.home() / '.askchat' / '.env'
 ENV_PATH = Path.home() / '.askchat' / 'envs'
 
+# Autocompletion
+# environment name completion
+class EnvNameCompletionType(click.ParamType):
+    name = "envname"
+    def shell_complete(self, ctx, param, incomplete):
+        return [
+            click.shell_completion.CompletionItem(path.stem) for path in ENV_PATH.glob(f"{incomplete}*.env")
+        ]
+# chat file completion
+class ChatFileCompletionType(click.ParamType):
+    name = "chatfile"
+    def shell_complete(self, ctx, param, incomplete):
+        return [
+            click.shell_completion.CompletionItem(path.stem) for path in CONFIG_PATH.glob(f"{incomplete}*.json")
+            if not path.name.startswith("_")
+        ]
+
+# common functions
 async def show_resp(chat):
     msg = ''
     async for char in chat.async_stream_responses(textonly=True):
diff --git a/askchat/askenv.py b/askchat/askenv.py
index 8d5bb1b..75e2787 100644
--- a/askchat/askenv.py
+++ b/askchat/askenv.py
@@ -1,5 +1,5 @@
 import click
-from askchat import write_config, ENV_PATH, MAIN_ENV_PATH
+from askchat import write_config, ENV_PATH, MAIN_ENV_PATH, EnvNameCompletionType
 from dotenv import set_key
 
 help_message = """Manage askchat environments.
@@ -38,14 +38,13 @@ def new(name, api_key, base_url, api_base, model):
     config_path = ENV_PATH / f'{name}.env'
     if config_path.exists():
         click.echo(f"Warning: Overwriting existing environment '{name}'.")
-        # use click.confirm to ask for user confirmation
         click.confirm("Do you want to continue?", abort=True)
     else:
         click.echo(f"Environment '{name}' created.")
     write_config(config_path, api_key, model, base_url, api_base)
 
 @cli.command()
-@click.argument('name', required=False)  # Make 'name' argument optional
+@click.argument('name', required=False, type=EnvNameCompletionType())
 @click.option('--default', is_flag=True, help='Delete the default environment configuration')
 def delete(name, default):
     """Delete an environment configuration."""
@@ -68,7 +67,7 @@ def delete(name, default):
             click.echo(f"Environment '{name}' not found.")
 
 @cli.command()
-@click.argument('name', required=False)
+@click.argument('name', required=False, type=EnvNameCompletionType())
 def show(name):
     """Print environment variables. Show default if no name is provided."""
     config_path = ENV_PATH / f'{name}.env' if name else MAIN_ENV_PATH
@@ -88,13 +87,16 @@ def save(name):
     if MAIN_ENV_PATH.exists():
         content = MAIN_ENV_PATH.read_text()
         config_path = ENV_PATH / f'{name}.env'
+        if config_path.exists():
+            click.echo(f"Warning: Overwriting existing environment '{name}'.")
+            click.confirm("Do you want to continue?", abort=True)
         config_path.write_text(content)
         click.echo(f"Environment '{name}' saved.")
     else:
         click.echo("No active environment to save.")
 
 @cli.command()
-@click.argument('name')
+@click.argument('name', type=EnvNameCompletionType())
 def use(name):
     """Activate an environment by replacing the .env file."""
     config_path = ENV_PATH / f'{name}.env'
@@ -110,7 +112,7 @@ def use(name):
 @click.option('-b', '--base-url', help='Base URL of the API (without suffix `/v1`)')
 @click.option('--api-base', help='Base URL of the API (with suffix `/v1`)')
 @click.option('-m', '--model', help='Model name')
-@click.argument('name', required=False)
+@click.argument('name', required=False, type=EnvNameCompletionType())
 def config(name, api_key, base_url, api_base, model):
     """Update default .env values."""
     if not any([api_key, base_url, api_base, model]):
diff --git a/askchat/cli.py b/askchat/cli.py
index 5b605ec..0addb1f 100644
--- a/askchat/cli.py
+++ b/askchat/cli.py
@@ -11,6 +11,7 @@
       show_resp, write_config
     , ENV_PATH, MAIN_ENV_PATH
     , CONFIG_PATH, CONFIG_FILE
+    , EnvNameCompletionType, ChatFileCompletionType
 )
 
 # Version and Config Path
@@ -86,13 +87,6 @@ def save_chat_callback(ctx, param, value):
         click.echo("No last conversation to save.")
     ctx.exit()
 
-class DeleteChatCompletionType(click.ParamType):
-    name = "deletechat"
-    def shell_complete(self, ctx, param, incomplete):
-        return [
-            click.shell_completion.CompletionItem(path.stem) for path in CONFIG_PATH.glob(f"{incomplete}*.json")
-        ]
-
 def delete_chat_callback(ctx, param, value):
     if not value:
         return
@@ -112,14 +106,6 @@ def list_chats_callback(ctx, param, value):
             click.echo(f" - {file.stem}")
     ctx.exit()
 
-# Custom type for environment name completion
-class EnvNameCompletionType(click.ParamType):
-    name = "envname"
-    def shell_complete(self, ctx, param, incomplete):
-        return [
-            click.shell_completion.CompletionItem(path.stem) for path in ENV_PATH.glob(f"{incomplete}*.env")
-        ]
-    
 # callback function for --use-env option
 def use_env_callback(ctx, param, value):
     if not value:
@@ -153,7 +139,7 @@ def cli():
 # Handling chat history
 @click.option('-p', '--print', is_flag=True, help='Print the last conversation or a specific conversation')
 @click.option('-s', '--save', callback=save_chat_callback, expose_value=False, help='Save the conversation to a file')
-@click.option('-d', '--delete', type=DeleteChatCompletionType(), callback=delete_chat_callback, expose_value=False, help='Delete the conversation from a file')
+@click.option('-d', '--delete', type=ChatFileCompletionType(), callback=delete_chat_callback, expose_value=False, help='Delete the conversation from a file')
 @click.option('--list', is_flag=True, callback=list_chats_callback, expose_value=False, help='List all the conversation files')
 # Other options
 @click.option('--generate-config', is_flag=True, callback=generate_config_callback, expose_value=False, help='Generate a configuration file by environment table')
diff --git a/setup.py b/setup.py
index b3768dc..3869cf4 100644
--- a/setup.py
+++ b/setup.py
@@ -4,12 +4,12 @@
 
 from setuptools import setup, find_packages
 
-VERSION = '1.1.0'
+VERSION = '1.1.1'
 
 with open('README.md') as readme_file:
     readme = readme_file.read()
 
-requirements = ['chattool>=3.1.1', "python-dotenv>=0.17.0", 'Click>=8.0']
+requirements = ['chattool>=3.1.2', "python-dotenv>=0.17.0", 'Click>=8.0']
 
 test_requirements = ['pytest>=3']
 
diff --git a/tests/test_askenv.py b/tests/test_askenv.py
index 0c65aa2..30d75b1 100644
--- a/tests/test_askenv.py
+++ b/tests/test_askenv.py
@@ -42,7 +42,6 @@ def test_list_initially_empty(runner, setup_env):
     result = runner.invoke(cli, ["list"])
     assert env_name in result.output
 
-
 def test_delete_environment(runner, setup_env):
     """Test deleting an environment configuration."""
     env_name, config_path = setup_env