set custom path for custom agent packages

instrukt/agent/loading.py

@@ -26,6 +26,7 @@
 import inspect
 import os
 from typing import Generator, Optional, Type, Any
+from xdg import BaseDirectory  # type: ignore
 
 from ..context import Context
 from ..errors import AgentError
@@ -36,7 +37,8 @@
 from pathlib import Path
 
 log = logging.getLogger(__name__)
-AGENT_MODULES_PATH = Path(__file__).parent.parent / "agent_modules"
+AGENT_MODULES_PATHS = [Path(__file__).parent.parent / "agent_modules",
+                      Path(BaseDirectory.save_data_path("instrukt/agents"))]
 MODULE_ENTRY_POINT = "main.py"
 MODULE_MANIFEST = "manifest.json"
 IMPL_CLS = InstruktAgent
@@ -48,26 +50,26 @@
 class ModuleManager:
 
     @staticmethod
-    def discover() -> Generator[str, None, None]:
+    def discover(paths: list[str]) -> Generator[str, None, None]:
         """ Discover available agent modules. Agent modules must be python packages.
 
-        
         If you started the agent module with a dot(.) or underscore(_) it will not be
         discoverable. Useful during development
         """
-        for _, name, ispkg in pkgutil.iter_modules([str(AGENT_MODULES_PATH)]):
-            if not ispkg:
-                raise ValueError(f"Agent module <{name}> must be a package")
+        for path in paths:
+            for _, name, ispkg in pkgutil.iter_modules([str(path)]):
+                if not ispkg:
+                    raise ValueError(f"Agent module <{name}> must be a package")
 
-            if name.startswith(".") or name.startswith("_"):
-                continue
+                if name.startswith(".") or name.startswith("_"):
+                    continue
 
-            yield name
+                yield name
 
     @staticmethod
     def list_modules() -> Generator[str, None, None]:
         """List all available agent modules."""
-        yield from ModuleManager.discover()
+        yield from ModuleManager.discover(AGENT_MODULES_PATHS)
 
     @staticmethod
     def get_manifest(name: str) -> AgentManifest:
@@ -79,16 +81,16 @@ def get_manifest(name: str) -> AgentManifest:
         Raises:
             ValueError: If the agent module does not have a manifest.
         """
-        module_path = os.path.join(AGENT_MODULES_PATH, name)
-        manifest_path = os.path.join(AGENT_MODULES_PATH, name, MODULE_MANIFEST)
-        if not os.path.isfile(manifest_path):
-            raise ValueError(f"Agent module <{name}> does not have a manifest")
-
-        with open(manifest_path, "r") as manifest_file:
-            manifest = AgentManifest(**json.load(manifest_file))
-            assert manifest.name == name, f"Agent module <{name}> must have the \
-same name as the agent's python package. Got <{manifest.name}> instead."
-            return manifest
+        for path in AGENT_MODULES_PATHS:
+            module_path = os.path.join(path, name)
+            manifest_path = os.path.join(path, name, MODULE_MANIFEST)
+            if os.path.isfile(manifest_path):
+                with open(manifest_path, "r") as manifest_file:
+                    manifest = AgentManifest(**json.load(manifest_file))
+                    assert manifest.name == name, f"Agent module <{name}> must have the \
+    same name as the agent's python package. Got <{manifest.name}> instead."
+                    return manifest
+        raise ValueError(f"Agent module <{name}> does not have a manifest")
 
     @staticmethod
     def verify_module(name: str) -> Type[InstruktAgent]:
@@ -100,50 +102,45 @@ def verify_module(name: str) -> Type[InstruktAgent]:
         Raises:
             ValueError: If the module is not valid.
         """
+        for path in AGENT_MODULES_PATHS:
+            mod_path = os.path.join(path, name)
 
-        mod_path = os.path.join(AGENT_MODULES_PATH, name)
-
-        # must be a package
-        if not os.path.isdir(mod_path):
-            raise ValueError(f"Agent module <{name}> must be a package")
-
-        # the package must have a main.py file
-        entry_mod = os.path.join(mod_path, MODULE_ENTRY_POINT)
-        if not os.path.isfile(entry_mod):
-            raise ValueError(MSG_MAIN_NOT_FOUND.format(name=name,
-                                                       mod_entry=MODULE_ENTRY_POINT,
-                                                       impl_cls=IMPL_CLS))
-
-
-        # the main.py file must have a clas the implements InstruktAgent class
-        mod = importlib.import_module(f"instrukt.agent_modules.{name}.{MODULE_ENTRY_POINT[:-3]}")
-        for _, agent_cls in inspect.getmembers(mod, inspect.isclass):
-            if issubclass(agent_cls, InstruktAgent) and agent_cls is not InstruktAgent:
-
-                # get the name and description attribute from the subclass
-                agent_name = getattr(agent_cls, "name", None)
-                agent_description = getattr(agent_cls, "description", None)
-
-                assert agent_name is not None, f"Agent <{name}> must have a \
-                                                name attribute"
-                assert agent_description is not None, f"Agent <{name}> must \
-                                                        have a description attribute"
-
-                #agent name must be the same as agent package
-                pkg_ns_path = f"instrukt.agent_modules.{name}"
-                if agent_name != name:
-                    raise ValueError(f"Agent <{name}> must have the same name \
-                                    as the agent's python package <{pkg_ns_path}>. \
-                                    Got <{agent_name}> instead.")
+            # must be a package
+            if not os.path.isdir(mod_path):
+                continue
 
-                break
+            # the package must have a main.py file
+            entry_mod = os.path.join(mod_path, MODULE_ENTRY_POINT)
+            if not os.path.isfile(entry_mod):
+                continue
+
+            # the main.py file must have a class the implements InstruktAgent class
+            mod = importlib.import_module(f"{name}.{MODULE_ENTRY_POINT[:-3]}")
+            for _, agent_cls in inspect.getmembers(mod, inspect.isclass):
+                if issubclass(agent_cls, InstruktAgent) and agent_cls is not InstruktAgent:
+
+                    # get the name and description attribute from the subclass
+                    agent_name = getattr(agent_cls, "name", None)
+                    agent_description = getattr(agent_cls, "description", None)
+
+                    assert agent_name is not None, f"Agent <{name}> must have a \
+                                                    name attribute"
+                    assert agent_description is not None, f"Agent <{name}> must \
+                                                            have a description attribute"
+
+                    # agent name must be the same as agent package
+                    pkg_ns_path = f"{path}.{name}"
+                    if agent_name != name:
+                        raise ValueError(f"Agent <{name}> must have the same name \
+                                        as the agent's python package <{pkg_ns_path}>. \
+                                        Got <{agent_name}> instead.")
+
+                    return agent_cls
 
         else:
             raise ValueError(f"Agent <{name}> must contain a class that implements \
                                  implements {IMPL_CLS}.")
 
-        return agent_cls
-
 
 
 class AgentLoader:
@@ -159,7 +156,7 @@ def load_agent(name: str, ctx: Context) -> Optional[InstruktAgent]:
         # try to load agent commands first
         try:
             importlib.import_module(
-                "instrukt.agent_modules.{}.commands".format(name))
+                "{}.commands".format(name))
             msg = LogMessage.info("Loaded [b]{}[/] commands.".format(name))
             ctx.notify(msg)
         except ImportError:

instrukt/config.py

@@ -133,6 +133,8 @@ def customise_sources(cls, init_settings, env_settings,
     llm_errors_logdir: str = os.path.join(
         BaseDirectory.save_cache_path("instrukt", "llm_errors"))
 
+    custom_agents_path: str = BaseDirectory.save_data_path("instrukt/agents")
+
     # TUI SETTINGS
     interface: TUISettings = Field(default_factory=TUISettings)
 
@@ -197,6 +199,10 @@ def __init__(self, config_file: str = "instrukt.yml"):
         else:
             self.save_config()
 
+        # add custom agent path to sys.path
+        import sys
+        sys.path.append(self.config.custom_agents_path)
+
     @property
     def C(self) -> Settings:
         """Shortcut to self.config."""