Creating Toolkits

Toolkits allow you to make tools available to all UnifAI agents as a service. This section shows you how to create, configure, and run your toolkits using both the JavaScript/TypeScript and Python SDKs.

Overview

A toolkit is a collection of tools that you register and serve to agents dynamically. It requires a Toolkit API Key — which you can get for free from UnifAI.

Toolkit Creation and Configuration

Let's break down the process of creating and configuring a toolkit:

Initialize the Toolkit

JavaScript/TypeScript

import { Toolkit } from 'unifai-sdk';

const toolkit = new Toolkit({ apiKey: 'YOUR_TOOLKIT_API_KEY' });

Python

import unifai

toolkit = unifai.Toolkit(api_key='YOUR_TOOLKIT_API_KEY')

Rust

use unifai_sdk::toolkit::*;

let mut toolkit = ToolkitService::new("YOUR_TOOLKIT_API_KEY");

Update Toolkit Details

You can optionally update the toolkit's name and description:

JavaScript/TypeScript

await toolkit.updateToolkit({ 
    name: "EchoChamber", 
    description: "What's in, what's out." 
});

Python

await toolkit.update_toolkit(
    name="EchoChamber",
    description="What's in, what's out."
)

# Or using asyncio.run():
# asyncio.run(toolkit.update_toolkit(name="EchoChamber", description="What's in, what's out."))

Rust

service
    .update_info(ToolkitInfo {
        name: "EchoChamber".to_string(),
        description: "What's in, what's out.".to_string(),
    })
    .await
    .unwrap();

Register Action Handlers

Register actions that your toolkit will provide. The payloadDescription can be any string or dictionary that contains enough information for agents to understand the payload format. It acts like API documentation that agents can read to determine what parameters to use.

Note that payloadDescription doesn't have to be in a certain format, as long as agents can understand it as natural language and generate the correct payload. Think of it as the comments and docs for your API, agents read it and decide what parameters to use. In practice, using JSON schema is recommended to match the format of training data.

JavaScript/TypeScript

toolkit.action(
    {
        action: "echo",
        actionDescription: "Echo the message",
        payloadDescription: {
            content: {
              type: "string",
              description: "The message to echo"
            }
        }
    },
    async (ctx, payload) => {
        return ctx.result(`You are agent <${ctx.agentId}>, you said "${payload?.content}".`);
    }
);

Python

@toolkit.action(
    action="echo",
    action_description="Echo the message",
    payload_description={
        "content": {
            "type": "string",
            "description": "The message to echo"
        }
    },
)
async def echo(ctx: unifai.ActionContext, payload={}): # can be a sync function too
    return ctx.Result(f'You are agent <{ctx.agent_id}>, you said "{payload.get("content")}".')

Rust

use thiserror::Error;
use unifai_sdk::{
    serde::{self, Deserialize, Serialize},
    serde_json::json,
    tokio,
    toolkit::*,
};

struct EchoChamber;

#[derive(Serialize, Deserialize)]
#[serde(crate = "serde")]
struct EchoChamberArgs {
    pub content: String,
}

#[derive(Debug, Error)]
#[error("Echo error")]
struct EchoChamberError;

impl Action for EchoChamber {
    const NAME: &'static str = "echo";

    type Error = EchoChamberError;
    type Args = EchoChamberArgs;
    type Output = String;

    async fn definition(&self) -> ActionDefinition {
        ActionDefinition {
            description: "Echo the message".to_string(),
            payload: json!({
                "content": {
                    "type": "string",
                    "description": "The message to echo"
                }
            }),
            payment: None,
        }
    }

    async fn call(
        &self,
        ctx: ActionContext,
        params: ActionParams<Self::Args>,
    ) -> Result<ActionResult<Self::Output>, Self::Error> {
        let output = format!(
            "You are agent <${}>, you said \"{}\".",
            ctx.agent_id, params.payload.content
        );

        Ok(ActionResult {
            payload: output,
            payment: None,
        })
    }
}

toolkit.add_action(EchoChamber);

Start the Toolkit

Finally, start serving your toolkit:

JavaScript/TypeScript

await toolkit.run();

Python

await toolkit.run()

# Or using asyncio.run():
# asyncio.run(toolkit.run())

Rust

let runner = toolkit.start().await.unwrap();
let _ = runner.await.unwrap();

Now your toolkit can be discovered and used by all agents connected to UnifAI.

Complete Example

Here's the complete code putting it all together:

JavaScript/TypeScript

import { Toolkit } from 'unifai-sdk';

async function main() {
  const toolkit = new Toolkit({ apiKey: 'YOUR_TOOLKIT_API_KEY' });

  // Optionally update the toolkit details
  await toolkit.updateToolkit({ 
    name: "EchoChamber", 
    description: "What's in, what's out." 
  });

  // Register an action handler
  toolkit.action(
    {
      action: "echo",
      actionDescription: "Echo the message",
      payloadDescription: {
        content: { type: "string" }
      }
    },
    async (ctx, payload) => {
      return ctx.result(`You are agent <${ctx.agentId}>, you said "${payload?.content}".`);
    }
  );

  // Start serving the toolkit
  await toolkit.run();
}

main().catch(console.error);

Python

import unifai
import asyncio

toolkit = unifai.Toolkit(api_key='YOUR_TOOLKIT_API_KEY')

# Update the toolkit details
asyncio.run(toolkit.update_toolkit(name="EchoChamber", description="What's in, what's out."))

# Register an action handler
@toolkit.action(
    action="echo",
    action_description="Echo the message",
    payload_description={"content": {"type": "string"}},
)
async def echo(ctx: unifai.ActionContext, payload={}):
    return ctx.Result(f'You are agent <{ctx.agent_id}>, you said "{payload.get("content")}".')

# Start serving the toolkit
asyncio.run(toolkit.run())

Rust

use thiserror::Error;
use unifai_sdk::{
    serde::{self, Deserialize, Serialize},
    serde_json::json,
    tokio,
    toolkit::*,
};

struct EchoChamber;

#[derive(Serialize, Deserialize)]
#[serde(crate = "serde")]
struct EchoChamberArgs {
    pub content: String,
}

#[derive(Debug, Error)]
#[error("Echo error")]
struct EchoChamberError;

impl Action for EchoChamber {
    const NAME: &'static str = "echo";

    type Error = EchoChamberError;
    type Args = EchoChamberArgs;
    type Output = String;

    async fn definition(&self) -> ActionDefinition {
        ActionDefinition {
            description: "Echo the message".to_string(),
            payload: json!({
                "content": {
                    "type": "string",
                    "description": "The message to echo",
                    "required": true
                }
            }),
            payment: None,
        }
    }

    async fn call(
        &self,
        ctx: ActionContext,
        params: ActionParams<Self::Args>,
    ) -> Result<ActionResult<Self::Output>, Self::Error> {
        let output = format!(
            "You are agent <${}>, you said \"{}\".",
            ctx.agent_id, params.payload.content
        );

        Ok(ActionResult {
            payload: output,
            payment: None,
        })
    }
}

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt().init();

    let mut toolkit = ToolkitService::new("YOUR_TOOLKIT_API_KEY");

    let info = ToolkitInfo {
        name: "EchoChamber".to_string(),
        description: "What's in, what's out.".to_string(),
    };

    toolkit.update_info(info).await.unwrap();

    toolkit.add_action(EchoChamber);

    let runner = toolkit.start().await.unwrap();
    let _ = runner.await.unwrap();
}

Toolkit Examples

We open source our official toolkits at unifai-network/unifai-toolkits.