Fix chains as commands.

agixt/Chain.py

@@ -27,9 +27,8 @@
 
 
 class Chain:
-    def __init__(self, user=DEFAULT_USER, ApiClient=None):
+    def __init__(self, user=DEFAULT_USER):
         self.user = user
-        self.ApiClient = ApiClient
         self.user_id = get_user_id(self.user)
 
     def get_chain(self, chain_name):

agixt/Extensions.py

@@ -5,6 +5,13 @@
 import logging
 import inspect
 from Globals import getenv, DEFAULT_USER
+from MagicalAuth import get_user_id
+from agixtsdk import AGiXTSDK
+from DB import (
+    get_session,
+    Chain as ChainDB,
+    User,
+)
 
 logging.basicConfig(
     level=getenv("LOG_LEVEL"),
@@ -30,10 +37,15 @@ def __init__(
         self.conversation_name = conversation_name
         self.conversation_id = conversation_id
         self.agent_id = agent_id
-        self.ApiClient = ApiClient
+        self.ApiClient = (
+            ApiClient
+            if ApiClient
+            else AGiXTSDK(base_uri=getenv("API_URL"), api_key=api_key)
+        )
         self.api_key = api_key
-        self.commands = self.load_commands()
         self.user = user
+        self.user_id = get_user_id(self.user)
+        self.commands = self.load_commands()
         if agent_config != None:
             if "commands" not in self.agent_config:
                 self.agent_config["commands"] = {}
@@ -46,6 +58,11 @@ def __init__(
                 "commands": {},
             }
 
+    async def execute_chain(self, chain_name, user_input="", **kwargs):
+        return await self.ApiClient.run_chain(
+            chain_name=chain_name, user_input=user_input, chain_args=kwargs
+        )
+
     def get_available_commands(self):
         if self.commands == []:
             return []
@@ -82,6 +99,60 @@ def get_command_args(self, command_name: str):
                     return command["command_args"]
         return {}
 
+    def get_chains(self):
+        session = get_session()
+        user_data = session.query(User).filter(User.email == DEFAULT_USER).first()
+        global_chains = (
+            session.query(ChainDB).filter(ChainDB.user_id == user_data.id).all()
+        )
+        chains = session.query(ChainDB).filter(ChainDB.user_id == self.user_id).all()
+        chain_list = []
+        for chain in chains:
+            chain_list.append(chain.name)
+        for chain in global_chains:
+            chain_list.append(chain.name)
+        session.close()
+        return chain_list
+
+    def get_chain_args(self, chain_name):
+        skip_args = [
+            "command_list",
+            "context",
+            "COMMANDS",
+            "date",
+            "conversation_history",
+            "agent_name",
+            "working_directory",
+            "helper_agent_name",
+        ]
+        chain_data = self.get_chain(chain_name=chain_name)
+        steps = chain_data["steps"]
+        prompt_args = []
+        args = []
+        for step in steps:
+            try:
+                prompt = step["prompt"]
+                prompt_category = (
+                    prompt["category"] if "category" in prompt else "Default"
+                )
+                if "prompt_name" in prompt:
+                    args = self.ApiClient.get_prompt_args(
+                        prompt_name=prompt["prompt_name"],
+                        prompt_category=prompt_category,
+                    )
+                elif "command_name" in prompt:
+                    args = Extensions().get_command_args(
+                        command_name=prompt["command_name"]
+                    )
+                elif "chain_name" in prompt:
+                    args = self.get_chain_args(chain_name=prompt["chain_name"])
+                for arg in args:
+                    if arg not in prompt_args and arg not in skip_args:
+                        prompt_args.append(arg)
+            except Exception as e:
+                logging.error(f"Error getting chain args: {e}")
+        return prompt_args
+
     def load_commands(self):
         try:
             settings = self.agent_config["settings"]
@@ -102,7 +173,6 @@ def load_commands(self):
                         command_function,
                     ) in command_class.commands.items():
                         params = self.get_command_params(command_function)
-                        # Store the module along with the function name
                         commands.append(
                             (
                                 command_name,
@@ -111,6 +181,21 @@ def load_commands(self):
                                 params,
                             )
                         )
+        chains = self.get_chains()
+        for chain in chains:
+            chain_args = self.get_chain_args(chain)
+            commands.append(
+                (
+                    chain,
+                    self.execute_chain,
+                    "run_chain",
+                    {
+                        "chain_name": chain,
+                        "user_input": "",
+                        **{arg: "" for arg in chain_args},
+                    },
+                )
+            )
         return commands
 
     def get_extension_settings(self):
@@ -138,9 +223,12 @@ def find_command(self, command_name: str):
             if module.__name__ in DISABLED_EXTENSIONS:
                 continue
             if name == command_name:
-                command_function = getattr(module, function_name)
-                return command_function, module, params  # Updated return statement
-        return None, None, None  # Updated return statement
+                if isinstance(module, type):  # It's a class
+                    command_function = getattr(module, function_name)
+                    return command_function, module, params
+                else:  # It's a function (for chains)
+                    return module, None, params
+        return None, None, None
 
     def get_commands_list(self):
         self.commands = self.load_commands()
@@ -171,6 +259,9 @@ async def execute_command(self, command_name: str, command_args: dict = None):
             logging.error(f"Command {command_name} not found")
             return f"Command {command_name} not found"
 
+        if command_args is None:
+            command_args = {}
+
         for param in params:
             if param not in command_args:
                 if param != "self" and param != "kwargs":
@@ -179,17 +270,18 @@ async def execute_command(self, command_name: str, command_args: dict = None):
         for param in command_args:
             if param not in params:
                 del args[param]
-        # try:
-        output = await getattr(
-            module(
-                **injection_variables,
-            ),
-            command_function.__name__,
-        )(**args)
-        # except Exception as e:
-        #    output = f"Error: {str(e)}"
-        logging.info(f"Command Output: {output}")
-        return output
+
+        if module is None:  # It's a chain
+            return await command_function(
+                chain_name=command_name, user_input="", **args
+            )
+        else:  # It's a regular command
+            return await getattr(
+                module(
+                    **injection_variables,
+                ),
+                command_function.__name__,
+            )(**args)
 
     def get_command_params(self, func):
         params = {}

agixt/endpoints/Chain.py

@@ -127,10 +127,7 @@ async def get_chain_args(
 ):
     if is_admin(email=user, api_key=authorization) != True:
         raise HTTPException(status_code=403, detail="Access Denied")
-    ApiClient = get_api_client(authorization=authorization)
-    chain_args = Chain(user=user, ApiClient=ApiClient).get_chain_args(
-        chain_name=chain_name
-    )
+    chain_args = Chain(user=user).get_chain_args(chain_name=chain_name)
     return {"chain_args": chain_args}
 
 

agixt/extensions/agixt_actions.py

@@ -179,9 +179,6 @@ def __init__(self, **kwargs):
             "Plan Multistep Task": self.plan_multistep_task,
             "Replace init in File": self.replace_init_in_file,
         }
-        user = kwargs["user"] if "user" in kwargs else "user"
-        for chain in Chain(user=user).get_chains():
-            self.commands[chain] = self.run_chain
         self.command_name = (
             kwargs["command_name"] if "command_name" in kwargs else "Smart Prompt"
         )
@@ -433,34 +430,6 @@ async def create_task_chain(
             i += 1
         return chain_name
 
-    async def run_chain(
-        self, input_for_task: str = "", start_from_step_number: int = 1
-    ):
-        """
-        Run a chain
-
-        Args:
-        input_for_task (str): The input for the task
-
-        Returns:
-        str: The response from the chain
-        """
-        try:
-            step_number = int(start_from_step_number)
-        except:
-            step_number = 1
-        response = await self.ApiClient.run_chain(
-            chain_name=self.command_name,
-            user_input=input_for_task,
-            agent_name=self.agent_name,
-            all_responses=False,
-            from_step=step_number,
-            chain_args={
-                "conversation_name": self.conversation_name,
-            },
-        )
-        return response
-
     def parse_openapi(self, data):
         """
         Parse OpenAPI data to extract endpoints

agixt/version

@@ -1 +1 @@
-v1.6.15
+v1.6.16