Custom Actions

Basic example

New actions are registered with the @register_action decorator.

For example, an action that splits a comma-separate string into a list of strings:

from aijson import register_action

@register_action
def split_text(text: str) -> list[str]:
    """
    Split text by comma
    """
    return text.split(",")

Invocation example:

version: "0.1"
flow:
  ask:
    action: llm
    prompt: List some examples of {{ thing }}
  extract:
    action: split_text
    text:
      link: ask

The name of the action defaults to the function’s name (split_text). The docstring (Split text by comma) is automatically picked up as its description, and shown in the language server on-hover.

We encourage adding typehints to your code (str and list[str]). They will help you find errors faster by showing more informative errors when something goes wrong. In no way are they required.

Customization

To manually override the the function name, description, and other options, invoke register_action with keyword arguments, for example:

from aijson import register_action

@register_action(
    name="my_split_text",
)
def split_text(text: str) -> list[str]:
    """
    Split text by comma
    """
    return text.split(",")

version: "0.1"
flow:
  ask:
    action: llm
    prompt: List some examples of {{ thing }}
  extract:
    action: my_split_text
    text:
      link: ask

The above example will register the same action under the name my_split_text.

Asynchronous example

Actions may be synchronous or asynchronous.

For example, an asynchronous action that visits a webpage and returns its contents with the aiohttp library:

from aijson import register_action
import aiohttp

@register_action
async def get_url(url):
    """
    Return the contents of the webpage.
    """
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as response:
            return await response.text()

version: "0.1"
flow:
  get_page_contents:
    action: get_url
    url:
      var: url
  summarize:
    action: llm
    prompt: |
      Summarize the contents of {{ url }}:
      ```
      {{ get_page_contents }}
      ```

Running with Custom Actions

In the VSCode language server and preview, actions are recursively discovered and imported in the current working directory.

However, when running a flow in code, the custom actions must be imported before the flow is executed.

Instead of manually importing your actions, you may register an entrypoint for the module that contains them.

For example, using poetry, with your actions located in my_package.actions, include the following at the end of your pyproject.toml:

[tool.poetry.plugins."aijson"]
actions = "my_package.actions"

When you publish your package on PyPI with the above entrypoint, anyone who installs it with pip will have access to your custom actions.

See the aijson-utils for an example of such an actionpack.