diff --git a/src/rlm/__init__.py b/src/rlm/__init__.py index 3fb4117..494ebf0 100644 --- a/src/rlm/__init__.py +++ b/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__) diff --git a/src/rlm/skill.py b/src/rlm/skill.py new file mode 100644 index 0000000..260d260 --- /dev/null +++ b/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 diff --git a/src/rlm/tools/ipython.py b/src/rlm/tools/ipython.py index 4e43d73..73026e0 100644 --- a/src/rlm/tools/ipython.py +++ b/src/rlm/tools/ipython.py @@ -151,7 +151,7 @@ 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) @@ -159,27 +159,12 @@ def _inject_startup(self): import nest_asyncio nest_asyncio.apply() - -class _CallableModule(types.ModuleType): - # Make `await (...)` shorthand for `await .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)