| import importlib |
| import os |
| from pathlib import Path |
| from typing import IO, Any |
|
|
| import yaml |
|
|
|
|
| class Loader(yaml.SafeLoader): |
| """ |
| A custom YAML loader that adds support for various custom tags: |
| |
| - !include: includes a YAML file as a subdocument |
| - !prompt: returns a prompt class based on the specified string |
| - !tool: returns a tool class based on the specified string |
| - !env: returns the value of an environment variable |
| - !file: returns the contents of a file |
| """ |
| def __init__(self, stream: IO[Any]) -> None: |
| """ |
| Initializes a new instance of the Loader class. |
| :param stream: The stream to load YAML from. |
| :type stream: IOBase |
| """ |
| self._root = Path(stream.name).resolve().parent |
| super(Loader, self).__init__(stream) |
| self.add_constructor("!include", Loader.include) |
| self.add_constructor("!prompt", Loader.prompt) |
| self.add_constructor("!tool", Loader.tool) |
| self.add_constructor("!env", Loader.env) |
| self.add_constructor("!file", Loader.file) |
|
|
| def include(self, node: yaml.Node) -> Any: |
| """ |
| Loads a YAML file from a path relative to the current file. Use this tag to include other agent configs as plugins. |
| |
| :param node: The YAML node to be loaded. |
| :type node: yaml.Node |
| :return: The loaded YAML file. |
| :rtype: Any |
| """ |
| filename = Path(self.construct_scalar(node)) |
| if not filename.is_absolute(): |
| filename = self._root / filename |
| with open(filename, 'r') as f: |
| return yaml.load(f, Loader) |
|
|
| def prompt(self, node: yaml.Node) -> Any: |
| """ |
| Returns a PromptTemplate class based on the specified string. |
| |
| :param node: The YAML node representing the prompt string. |
| :type node: yaml.Node |
| :return: The prompt class. |
| :rtype: type |
| :raises AssertionError: If the resolved prompt class is not a subclass of PromptTemplate. |
| """ |
| from ..prompt import PromptTemplate, SimpleReactPrompt, ZeroShotReactPrompt |
| prompt = self.construct_scalar(node) |
| if '.' in prompt: |
| _path = prompt.split('.') |
| module = importlib.import_module('.'.join(_path[:-1])) |
| prompt_cls = getattr(module, _path[-1]) |
| else: |
| prompt_cls = eval(prompt) |
| assert issubclass(prompt_cls.__class__, PromptTemplate) |
| return prompt_cls |
|
|
| def tool(self, node: yaml.Node) -> Any: |
| """ |
| Loads a Custom BaseTool class from a path relative to the current file. |
| |
| :param node: The YAML node to be loaded. |
| :type node: yaml.Node |
| :return: The loaded BaseTool class. |
| :rtype: Any |
| """ |
| from ..tools import BaseTool, PythonSandBoxTool |
|
|
| tool = self.construct_scalar(node) |
| if '.' in tool: |
| _path = tool.split('.') |
| module = importlib.import_module('.'.join(_path[:-1])) |
| tool_cls = getattr(module, _path[-1]) |
| else: |
| tool_cls = eval(tool) |
| |
| assert issubclass(tool_cls, BaseTool) |
| return tool_cls |
|
|
| def env(self, node: yaml.Node) -> Any: |
| """ |
| Loads an environment variable from the current environment, defaults to an empty string if the variable is not set. |
| |
| :param node: The YAML node to be loaded. |
| :type node: yaml.Node |
| :return: The loaded environment variable. |
| :rtype: Any |
| """ |
| return os.environ.get(self.construct_scalar(node), "") |
|
|
| def file(self, node: yaml.Node) -> Any: |
| """ |
| Loads any readable file from a path relative to the current file. |
| |
| :param node: The YAML node to be loaded. |
| :type node: yaml.Node |
| :return: The loaded file. |
| :rtype: Any |
| """ |
| filename = Path(self.construct_scalar(node)) |
| if not filename.is_absolute(): |
| filename = self._root / filename |
| with open(filename, 'r') as f: |
| return f.read().strip() |
