Merge pull request #136 from tcdent/runtime-errors

Some friendly error messages for runtime errors in `agentstack run`

Author: tcdent

agentstack/cli/run.py

@@ -1,5 +1,6 @@
 from typing import Optional
 import sys
+import traceback
 from pathlib import Path
 import importlib.util
 from dotenv import load_dotenv
@@ -14,6 +15,62 @@
 MAIN_MODULE_NAME = "main"
 
 
+def _format_friendy_error_message(exception: Exception):
+    """
+    Projects will throw various errors, especially on first runs, so we catch
+    them here and print a more helpful message.
+
+    In order to prevent us from having to import all possible backend exceptions
+    we do string matching on the exception type and traceback contents.
+    """
+    # TODO These end up being pretty framework-specific; consider individual implementations.
+    COMMON_LLM_ENV_VARS = (
+        'OPENAI_API_KEY',
+        'ANTHROPIC_API_KEY',
+    )
+
+    name = exception.__class__.__name__
+    message = str(exception)
+    tracebacks = traceback.format_exception(type(exception), exception, exception.__traceback__)
+
+    match (name, message, tracebacks):
+        # The user doesn't have an environment variable set for the LLM provider.
+        case ('AuthenticationError', m, t) if 'litellm.AuthenticationError' in t[-1]:
+            variable_name = [k for k in COMMON_LLM_ENV_VARS if k in message] or ["correct"]
+            return (
+                "We were unable to connect to the LLM provider. "
+                f"Ensure your .env file has the {variable_name[0]} variable set."
+            )
+        # This happens when the LLM configured for an agent is invalid.
+        case ('BadRequestError', m, t) if 'LLM Provider NOT provided' in t[-1]:
+            return (
+                "An invalid LLM was configured for an agent. "
+                "Ensure the 'llm' attribute of the agent in the agents.yaml file is in the format <provider>/<model>."
+            )
+        # The user has not configured the correct agent name in the tasks.yaml file.
+        case ('KeyError', m, t) if 'self.tasks_config[task_name]["agent"]' in t[-2]:
+            return (
+                f"The agent {message} is not defined in your agents file. "
+                "Ensure the 'agent' fields in your tasks.yaml correspond to an entry in the agents.yaml file."
+            )
+        # The user does not have an agent defined in agents.yaml file, but it does
+        # exist in the entrypoint code.
+        case ('KeyError', m, t) if 'config=self.agents_config[' in t[-2]:
+            return (
+                f"The agent {message} is not defined in your agents file. "
+                "Ensure all agents referenced in your code are defined in the agents.yaml file."
+            )
+        # The user does not have a task defined in tasks.yaml file, but it does
+        # exist in the entrypoint code.
+        case ('KeyError', m, t) if 'config=self.tasks_config[' in t[-2]:
+            return (
+                f"The task {message} is not defined in your tasks. "
+                "Ensure all tasks referenced in your code are defined in the tasks.yaml file."
+            )
+        case (_, _, _):
+            return f"{name}: {message}, {tracebacks[-1]}"
+
+
 def _import_project_module(path: Path):
     """
     Import `main` from the project path.
@@ -32,7 +89,7 @@ def _import_project_module(path: Path):
     return project_module
 
 
-def run_project(command: str = 'run', cli_args: Optional[str] = None):
+def run_project(command: str = 'run', debug: bool = False, cli_args: Optional[str] = None):
     """Validate that the project is ready to run and then run it."""
     if conf.get_framework() not in frameworks.SUPPORTED_FRAMEWORKS:
         print(term_color(f"Framework {conf.get_framework()} is not supported by agentstack.", 'red'))
@@ -55,14 +112,18 @@ def run_project(command: str = 'run', cli_args: Optional[str] = None):
     load_dotenv(Path.home() / '.env')  # load the user's .env file
     load_dotenv(conf.PATH / '.env', override=True)  # load the project's .env file
 
-    # import src/main.py from the project path
+    # import src/main.py from the project path and run `command` from the project's main.py
     try:
+        print("Running your agent...")
         project_main = _import_project_module(conf.PATH)
+        getattr(project_main, command)()
     except ImportError as e:
         print(term_color(f"Failed to import project. Does '{MAIN_FILENAME}' exist?:\n{e}", 'red'))
         sys.exit(1)
-
-    # run `command` from the project's main.py
-    # TODO try/except this and print detailed information with a --debug flag
-    print("Running your agent...")
-    return getattr(project_main, command)()
+    except Exception as exception:
+        if debug:
+            raise exception
+        print(term_color("\nAn error occurred while running your project:\n", 'red'))
+        print(_format_friendy_error_message(exception))
+        print(term_color("\nRun `agentstack run --debug` for a full traceback.", 'blue'))
+        sys.exit(1)

agentstack/main.py

@@ -24,6 +24,12 @@ def main():
         help="Path to the project directory, defaults to current working directory",
         dest="project_path",
     )
+    global_parser.add_argument(
+        "--debug",
+        help="Print more information when an error occurs",
+        dest="debug",
+        action="store_true",
+    )
 
     parser = argparse.ArgumentParser(
         parents=[global_parser], description="AgentStack CLI - The easiest way to build an agent application"
@@ -157,7 +163,7 @@ def main():
     elif args.command in ["init", "i"]:
         init_project_builder(args.slug_name, args.template, args.wizard)
     elif args.command in ["run", "r"]:
-        run_project(command=args.function, cli_args=extra_args)
+        run_project(command=args.function, debug=args.debug, cli_args=extra_args)
     elif args.command in ['generate', 'g']:
         if args.generate_command in ['agent', 'a']:
             if not args.llm: