Skip to content

Commit

Permalink
add autocompletion for askenv
Browse files Browse the repository at this point in the history
  • Loading branch information
RexWzh committed Mar 12, 2024
1 parent 30542bb commit 0b7d7d1
Show file tree
Hide file tree
Showing 5 changed files with 32 additions and 26 deletions.
21 changes: 20 additions & 1 deletion askchat/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,17 +2,36 @@

__author__ = """Rex Wang"""
__email__ = '[email protected]'
__version__ = '1.1.0'
__version__ = '1.1.1'

import asyncio
from pathlib import Path
import click

# Main environment file
CONFIG_PATH = Path.home() / ".askchat"
CONFIG_FILE = CONFIG_PATH / ".env"
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):
Expand Down
14 changes: 8 additions & 6 deletions askchat/askenv.py
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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."""
Expand All @@ -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
Expand All @@ -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'
Expand All @@ -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]):
Expand Down
18 changes: 2 additions & 16 deletions askchat/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
show_resp, write_config
, ENV_PATH, MAIN_ENV_PATH
, CONFIG_PATH, CONFIG_FILE
, EnvNameCompletionType, ChatFileCompletionType
)

# Version and Config Path
Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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')
Expand Down
4 changes: 2 additions & 2 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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']

Expand Down
1 change: 0 additions & 1 deletion tests/test_askenv.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 0b7d7d1

Please sign in to comment.