skills: opt-in callable_module helper for await <skill>(...)

src/rlm/__init__.py

@@ -2,6 +2,12 @@
 
 from rlm.api import run
 from rlm.engine import RLMEngine
+from rlm.skill import callable_module
 from rlm.types import RLMMetrics, RLMResult
 
-__all__ = ["run", "RLMEngine", "RLMMetrics", "RLMResult"]
+__all__ = ["callable_module", "run", "RLMEngine", "RLMMetrics", "RLMResult"]
+
+# Opt the rlm module itself into the callable-shorthand helper we ship
+# for skill authors, so `await rlm('sub-task')` works identically to
+# `await rlm.run('sub-task')` inside an IPython cell.
+callable_module(__name__)

src/rlm/skill.py

@@ -0,0 +1,28 @@
+"""Public helpers for building rlm skills."""
+
+from __future__ import annotations
+
+import sys
+import types
+
+
+def callable_module(name: str) -> None:
+    """Make a skill module directly awaitable.
+
+    Call from the skill package's ``__init__.py`` after ``run`` is bound
+    at module scope::
+
+        from .edit import PARAMETERS, main, run
+        from rlm.skill import callable_module
+        callable_module(__name__)
+
+    After this runs, ``await edit(...)`` is equivalent to
+    ``await edit.run(...)``; ``edit.run`` and ``edit.PARAMETERS`` stay
+    accessible unchanged.
+    """
+
+    class _CallableSkill(types.ModuleType):
+        async def __call__(self, *args, **kwargs):
+            return await self.run(*args, **kwargs)
+
+    sys.modules[name].__class__ = _CallableSkill

src/rlm/tools/ipython.py

@@ -151,35 +151,20 @@ def _inject_startup(self):
         installed_skills = get_installed_skills()
 
         setup_code = f"""\
-import os, sys, types
+import os
 os.chdir({self.cwd!r})
 os.environ['RLM_SESSION_DIR'] = {session_dir!r} or ''
 os.environ['RLM_DEPTH'] = str({depth!r} + 1)
 
 import nest_asyncio
 nest_asyncio.apply()
 
-
-class _CallableModule(types.ModuleType):
-    # Make `await <skill>(...)` shorthand for `await <skill>.run(...)`.
-    # __call__ is looked up on the type, not the instance, so the
-    # override has to live on the class.
-    async def __call__(self, *args, **kwargs):
-        return await self.run(*args, **kwargs)
-
-
-def _wrap_callable(mod):
-    wrapped = _CallableModule(mod.__name__)
-    wrapped.__dict__.update(mod.__dict__)
-    sys.modules[mod.__name__] = wrapped
-    return wrapped
-
-
 for _name in {installed_skills!r}:
-    globals()[_name] = _wrap_callable(__import__(_name))
+    globals()[_name] = __import__(_name)
 
 if {allow_recursion!r}:
-    globals()['rlm'] = _wrap_callable(__import__('rlm'))
+    import rlm
+    globals()['rlm'] = rlm
 """
         self._execute_silent(setup_code)