add help and descriptions to cli

ollama/cmd/cli.py

@@ -1,32 +1,76 @@
 import os
 import sys
-from argparse import ArgumentParser
+from argparse import ArgumentParser, HelpFormatter, PARSER
 from yaspin import yaspin
 
 from ollama import model, engine
 from ollama.cmd import server
 
 
+class CustomHelpFormatter(HelpFormatter):
+    """
+    This class is used to customize the way the argparse help text is displayed.
+    We specifically override the _format_action method to exclude the line that
+    shows all the subparser command options in the help text. This line is typically
+    in the form "{serve,models,pull,run}".
+    """
+
+    def _format_action(self, action):
+        # get the original help text
+        parts = super()._format_action(action)
+        if action.nargs == PARSER:
+            # remove the unwanted first line
+            parts = "\n".join(parts.split("\n")[1:])
+        return parts
+
+
 def main():
-    parser = ArgumentParser()
+    parser = ArgumentParser(
+        description='Ollama: Run any large language model on any machine.',
+        formatter_class=CustomHelpFormatter,
+    )
 
     # create models home if it doesn't exist
     os.makedirs(model.models_home, exist_ok=True)
 
-    subparsers = parser.add_subparsers()
+    subparsers = parser.add_subparsers(
+        title='commands',
+    )
 
-    server.set_parser(subparsers.add_parser("serve"))
+    server.set_parser(
+        subparsers.add_parser(
+            "serve",
+            description="Start a persistent server to interact with models via the API.",
+            help="Start a persistent server to interact with models via the API.",
+        )
+    )
 
-    list_parser = subparsers.add_parser("models")
+    list_parser = subparsers.add_parser(
+        "models",
+        description="List all available models stored locally.",
+        help="List all available models stored locally.",
+    )
     list_parser.set_defaults(fn=list_models)
 
-    pull_parser = subparsers.add_parser("pull")
-    pull_parser.add_argument("model")
+    pull_parser = subparsers.add_parser(
+        "pull",
+        description="Download a specified model from a remote source.",
+        help="Download a specified model from a remote source. Usage: pull [model]",
+    )
+    pull_parser.add_argument("model", help="Name of the model to download.")
     pull_parser.set_defaults(fn=pull)
 
-    run_parser = subparsers.add_parser("run")
-    run_parser.add_argument("model")
-    run_parser.add_argument("prompt", nargs="?")
+    run_parser = subparsers.add_parser(
+        "run",
+        description="Run a model and submit prompts.",
+        help="Run a model and submit prompts. Usage: run [model] [prompt]",
+    )
+    run_parser.add_argument("model", help="Name of the model to run.")
+    run_parser.add_argument(
+        "prompt",
+        nargs="?",
+        help="Optional prompt for the model, interactive mode enabled when not specified.",
+    )
     run_parser.set_defaults(fn=run)
 
     args = parser.parse_args()