Fix wizard!

Author: tcdent

agentstack/cli/wizard.py

@@ -5,28 +5,78 @@
 import webbrowser
 from art import text2art
 from agentstack import log
+from agentstack.frameworks import SUPPORTED_FRAMEWORKS
 from agentstack.utils import open_json_file, is_snake_case
 from agentstack.cli import welcome_message, get_validated_input
+from agentstack.cli.cli import PREFERRED_MODELS
+from agentstack._tools import get_all_tools, get_all_tool_names
 from agentstack.proj_templates import TemplateConfig
 
 
-def run_wizard(slug_name: str) -> TemplateConfig:
-    raise NotImplementedError("TODO wizard functionality needs to be migrated")
+class WizardData(dict):
+    def to_template_config(self) -> TemplateConfig:
+        agents = []
+        for agent in self['design']['agents']:
+            agents.append(TemplateConfig.Agent(**{
+                'name': agent['name'],
+                'role': agent['role'],
+                'goal': agent['goal'],
+                'backstory': agent['backstory'],
+                'llm': agent['model'],
+            }))
+        
+        tasks = []
+        for task in self['design']['tasks']:
+            tasks.append(TemplateConfig.Task(**{
+                'name': task['name'],
+                'description': task['description'],
+                'expected_output': task['expected_output'],
+                'agent': task['agent'],
+            }))
+        
+        tools = []
+        for tool in self['tools']:
+            tools.append(TemplateConfig.Tool(**{
+                'name': tool,
+                'agents': [agent.name for agent in agents],  # all agents
+            }))
+        
+        return TemplateConfig(
+            name=self['project']['name'],
+            description=self['project']['description'],
+            template_version=4,
+            framework=self['framework'],
+            method='sequential',
+            manager_agent=None,
+            agents=agents,
+            tasks=tasks,
+            tools=tools,
+            graph=[],
+            inputs={},
+        )
 
+
+def run_wizard(slug_name: str) -> TemplateConfig:
     project_details = ask_project_details(slug_name)
     welcome_message()
     framework = ask_framework()
     design = ask_design()
     tools = ask_tools()
-     # TODO return TemplateConfig object
+    
+    wizard_data = WizardData({
+        'project': project_details,
+        'framework': framework,
+        'design': design,
+        'tools': tools,
+    })
+    return wizard_data.to_template_config()
 
 
 def ask_framework() -> str:
-    framework = "CrewAI"
-    # framework = inquirer.list_input(
-    #     message="What agent framework do you want to use?",
-    #     choices=["CrewAI", "Autogen", "LiteLLM", "Learn what these are (link)"],
-    # )
+    framework = inquirer.list_input(
+        message="What agent framework do you want to use?",
+        choices=SUPPORTED_FRAMEWORKS,
+    )
     #
     # if framework == "Learn what these are (link)":
     #     webbrowser.open("https://youtu.be/xvFZjo5PgG0")
@@ -42,8 +92,7 @@ def ask_framework() -> str:
     #         choices=["CrewAI", "Autogen", "LiteLLM"],
     #     )
 
-    log.success("Congrats! Your project is ready to go! Quickly add features now or skip to do it later.\n\n")
-
+    #log.success("Congrats! Your project is ready to go! Quickly add features now or skip to do it later.\n\n")
     return framework
 
 
@@ -157,20 +206,28 @@ def ask_tools() -> list:
     tools_to_add = []
 
     adding_tools = True
-    script_dir = os.path.dirname(os.path.abspath(__file__))
-    tools_json_path = os.path.join(script_dir, '..', 'tools', 'tools.json')
-
-    # Load the JSON data
-    tools_data = open_json_file(tools_json_path)
+    tool_configs = get_all_tools()
 
     while adding_tools:
+        tool_categories = []
+        for tool_config in tool_configs:
+            if tool_config.category not in tool_categories:
+                tool_categories.append(tool_config.category)
+        
         tool_type = inquirer.list_input(
             message="What category tool do you want to add?",
-            choices=list(tools_data.keys()) + ["~~ Stop adding tools ~~"],
+            choices=tool_categories + ["~~ Stop adding tools ~~"],
         )
 
-        tools_in_cat = [f"{t['name']} - {t['url']}" for t in tools_data[tool_type] if t not in tools_to_add]
-        tool_selection = inquirer.list_input(message="Select your tool", choices=tools_in_cat)
+        tools_in_cat = []
+        for tool_config in tool_configs:
+            if tool_config.category == tool_type:
+                tools_in_cat.append(tool_config)
+        
+        tool_selection = inquirer.list_input(
+            message="Select your tool", 
+            choices=[f"{t.name} - {t.url}" for t in tools_in_cat if t not in tools_to_add], 
+        )
 
         tools_to_add.append(tool_selection.split(' - ')[0])
 

agentstack/proj_templates.py

@@ -168,8 +168,8 @@ class TemplateConfig(pydantic.BaseModel):
         A list of tools used by the project.
     graph: list[list[TemplateConfig.Node]]
         A list of graph relationships. Each edge must have exactly 2 nodes.
-    inputs: list[str]
-        A list of inputs used by the project.
+    inputs: dict[str, str]
+        Key/value pairs of inputs used by the project.
     """
 
     class Agent(pydantic.BaseModel):