Custom Toolkits

1from typing import List
2
3from kern.agent import Agent
4from kern.tools import Toolkit
5from kern.utils.log import logger
6
7class ShellTools(Toolkit):
8    def __init__(self, working_directory: str = "/", **kwargs):
9        self.working_directory = working_directory
10
11        tools = [
12            self.run_shell_command,
13        ]
14
15        super().__init__(name="shell_tools", tools=tools, **kwargs)
16    
17    def list_files(self, directory: str):
18        """
19        List the files in the given directory.
20
21        Args:
22            directory (str): The directory to list the files from.
23        Returns:
24            str: The list of files in the directory.
25        """
26        import os
27
28        # List files relative to the toolkit's working_directory
29        path = os.path.join(self.working_directory, directory)
30        try:
31            files = os.listdir(path)
32            return "\n".join(files)
33        except Exception as e:
34            logger.warning(f"Failed to list files in {path}: {e}")
35            return f"Error: {e}"
36        return os.listdir(directory)
37
38    def run_shell_command(self, args: List[str], tail: int = 100) -> str:
39        """
40        Runs a shell command and returns the output or error.
41
42        Args:
43            args (List[str]): The command to run as a list of strings.
44            tail (int): The number of lines to return from the output.
45        Returns:
46            str: The output of the command.
47        """
48        import subprocess
49
50        logger.info(f"Running shell command: {args}")
51        try:
52            logger.info(f"Running shell command: {args}")
53            result = subprocess.run(args, capture_output=True, text=True, cwd=self.working_directory)
54            logger.debug(f"Result: {result}")
55            logger.debug(f"Return code: {result.returncode}")
56            if result.returncode != 0:
57                return f"Error: {result.stderr}"
58            # return only the last n lines of the output
59            return "\n".join(result.stdout.split("\n")[-tail:])
60        except Exception as e:
61            logger.warning(f"Failed to run shell command: {e}")
62            return f"Error: {e}"
63
64agent = Agent(tools=[ShellTools()], markdown=True)
65agent.print_response("List all the files in my home directory.")

1from typing import Any, Dict
2
3from kern.agent import Agent
4from kern.tools import Toolkit
5
6try:
7    import httpx
8except ImportError:
9    raise ImportError("`httpx` not installed. Run `uv pip install httpx`")
10
11
12class APITools(Toolkit):
13    def __init__(self, base_url: str, timeout: float = 30.0, **kwargs):
14        self.base_url = base_url
15        self.timeout = timeout
16
17        # Sync tools for agent.run() and agent.print_response()
18        tools = [
19            self.fetch_data,
20            self.post_data,
21        ]
22
23        # Async tools for agent.arun() and agent.aprint_response()
24        # Format: (async_method, "tool_name")
25        async_tools = [
26            (self.afetch_data, "fetch_data"),
27            (self.apost_data, "post_data"),
28        ]
29
30        super().__init__(name="api_tools", tools=tools, async_tools=async_tools, **kwargs)
31
32    # Sync methods
33    def fetch_data(self, endpoint: str) -> Dict[str, Any]:
34        """
35        Fetch data from an API endpoint.
36
37        Args:
38            endpoint: The API endpoint to fetch data from (e.g., "/users/123")
39        Returns:
40            The JSON response from the API
41        """
42        url = f"{self.base_url}{endpoint}"
43        with httpx.Client(timeout=self.timeout) as client:
44            response = client.get(url)
45            response.raise_for_status()
46            return response.json()
47
48    def post_data(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
49        """
50        Post data to an API endpoint.
51
52        Args:
53            endpoint: The API endpoint to post data to
54            data: The data to post as JSON
55        Returns:
56            The JSON response from the API
57        """
58        url = f"{self.base_url}{endpoint}"
59        with httpx.Client(timeout=self.timeout) as client:
60            response = client.post(url, json=data)
61            response.raise_for_status()
62            return response.json()
63
64    # Async methods (used automatically in async contexts)
65    async def afetch_data(self, endpoint: str) -> Dict[str, Any]:
66        """
67        Fetch data from an API endpoint asynchronously.
68
69        Args:
70            endpoint: The API endpoint to fetch data from (e.g., "/users/123")
71        Returns:
72            The JSON response from the API
73        """
74        url = f"{self.base_url}{endpoint}"
75        async with httpx.AsyncClient(timeout=self.timeout) as client:
76            response = await client.get(url)
77            response.raise_for_status()
78            return response.json()
79
80    async def apost_data(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
81        """
82        Post data to an API endpoint asynchronously.
83
84        Args:
85            endpoint: The API endpoint to post data to
86            data: The data to post as JSON
87        Returns:
88            The JSON response from the API
89        """
90        url = f"{self.base_url}{endpoint}"
91        async with httpx.AsyncClient(timeout=self.timeout) as client:
92            response = await client.post(url, json=data)
93            response.raise_for_status()
94            return response.json()
95
96# Create the agent with the toolkit (using JSONPlaceholder - a free fake API for testing)
97agent = Agent(tools=[APITools(base_url="https://jsonplaceholder.typicode.com")], markdown=True)
98
99# Sync usage - uses fetch_data
100agent.print_response("Fetch the user with ID 1")
101
102# Async usage - uses afetch_data automatically
103import asyncio
104asyncio.run(agent.aprint_response("Fetch the post with ID 1"))