Enhance CLI with subcommands and provider logic

- Introduces `version` subcommand to display the app version
- Prevents simultaneous use of subcommands and other flags
- Implements `handle_subcommand` to manage subcommand execution
- Adds type annotations for CLI arguments using `Annotated`
- Modifies `get_api_key` to return `None` when no API key is found
- Updates `get_ai_provider` to dynamically set the model for the ollama provider
- Ensures a model name is set for remote providers via environment variable or CLI flag

Author: yankeexe

ai_commit/__init__.py

@@ -1,12 +1,14 @@
 import argparse
 from dataclasses import dataclass
+from typing import Annotated
 
 
 @dataclass
 class CLIArgs:
-    remote: bool
-    debug: bool
-    model: str | None = None
+    remote: Annotated[bool, "Use remote model for commit generation"]
+    debug: Annotated[bool, "Run the CLI in debug mode"]
+    model: Annotated[str | None, "Model to use for commit generation"] = None
+    command: Annotated[str | None, "Sub-command passed to the CLI"] = None
 
 
 parser = argparse.ArgumentParser(
@@ -34,5 +36,15 @@ class CLIArgs:
     "-r", "--remote", help=remote_model_help, default=False, action="store_true"
 )
 
+
+version_parser = parser.add_subparsers(title="version", dest="command")
+version_parser.add_parser("version", help="Show app version")
 raw_args = parser.parse_args()
-cli_args = CLIArgs(remote=raw_args.remote, debug=raw_args.debug, model=raw_args.model)
+
+
+cli_args = CLIArgs(
+    remote=raw_args.remote,
+    debug=raw_args.debug,
+    model=raw_args.model,
+    command=raw_args.command,
+)

ai_commit/app.py

@@ -4,10 +4,11 @@
 import subprocess
 import sys
 import tempfile
+from importlib.metadata import PackageNotFoundError, version as package_version
 
 from openai import OpenAI
 
-from ai_commit import cli_args as args
+from ai_commit import CLIArgs, cli_args as args
 from ai_commit.prompts import get_system_prompt
 from ai_commit.providers import get_ai_provider, provider_names
 
@@ -19,7 +20,7 @@
 }
 
 
-def get_api_key() -> str:
+def get_api_key() -> str | None:
     """Check if API keys are set in environment variables for remote models.
     If the model is run locally (via Ollama), then the 'ollama' key is returned back.
 
@@ -108,7 +109,7 @@ def generate_commit_message(staged_changes: str, regenerate: bool = False) -> st
 
         return commit_message
     except Exception as e:
-        print(f"❌ Error generating commit message: {str(e)}")
+        print(f"❌ Error generating commit message:\n{str(e)}")
         sys.exit(1)
 
 
@@ -163,7 +164,6 @@ def handle_editing(commit_message: str):
 def interaction_loop(staged_changes: str):
     commit_message = generate_commit_message(staged_changes)
     while True:
-
         action = input(
             "\n\nProceed to commit? [y(yes) | n[no] | r(regenerate) | e(edit)] "
         )
@@ -197,8 +197,29 @@ def interaction_loop(staged_changes: str):
                 break
 
 
+def get_version() -> str:
+    try:
+        version = package_version("ai-gen-commit")
+    except PackageNotFoundError:
+        version = "unknown"
+    return version
+
+
+def handle_subcommand(cli_args: CLIArgs):
+    if cli_args.command and (cli_args.model or cli_args.debug or cli_args.remote):
+        print(
+            f"❌ Error: cannot use subcommand {cli_args.command} and other flags together"
+        )
+        sys.exit(1)
+
+    if cli_args.command == "version":
+        print(get_version())
+        sys.exit(0)
+
+
 def run():
     try:
+        handle_subcommand(args)
         run_command(commands["is_git_repo"])
         staged_changes = run_command(commands["get_stashed_changes"])
 

ai_commit/providers.py

@@ -9,14 +9,14 @@
 @dataclass
 class Provider:
     name: str
-    model: str
+    model: str | None
     base_url: str
 
 
 providers_mapping = {
     "ollama": Provider(
         name="ollama",
-        model=get_model(),
+        model=None,
         base_url=parse_host(os.getenv("OLLAMA_HOST")) + "/v1",
     ),
     "openai": Provider(name="openai", model="gpt-4o", base_url=""),
@@ -61,10 +61,12 @@ class Provider:
 
 def get_ai_provider() -> Provider | None:
     if not args.remote:
-        return providers_mapping["ollama"]
+        provider = providers_mapping["ollama"]
+        if provider.model is None:
+            provider.model = get_model()
+        return provider
 
     ai_provider = os.environ.get("AI_COMMIT_PROVIDER")
-    ai_provider = ai_provider.lower() if ai_provider else None
     if not ai_provider:
         print(
             f"""
@@ -84,4 +86,18 @@ def get_ai_provider() -> Provider | None:
         )
         sys.exit(1)
 
-    return providers_mapping.get(ai_provider)
+    provider = providers_mapping.get(ai_provider.lower())
+    if not provider.model:
+        print(
+            f"""
+❌ No model name set for the provider: {provider.name}
+
+Set model name using:
+    export AI_COMMIT_MODEL=<model-from-provider>
+OR 
+    aic -m <model-from-provider>
+                    """
+        )
+        sys.exit(1)
+
+    return providers_mapping.get(ai_provider.lower())