Files
NexQuant/rdagent/utils/agent/tpl.py
T
Linlang faa2fb03ad CI checks that can be automatically repaired (#119)
* fix isort & black & toml-sort & sphinx error

* fix ci error

* fix ci error

* add comments

* Update Makefile

* change sphinx build command

* add auto-lint

* add black args

* format with black

* Auto Linting document

* fix ci error

---------

Co-authored-by: you-n-g <you-n-g@users.noreply.github.com>
Co-authored-by: Young <afe.young@gmail.com>
2024-07-26 12:12:16 +08:00

64 lines
2.0 KiB
Python

"""
Here are some infrastruture to build a agent
The motivation of tempalte and AgentOutput Design
"""
import inspect
from pathlib import Path
from typing import Any
import yaml
from jinja2 import Environment, StrictUndefined
from rdagent.core.utils import SingletonBaseClass
DIRNAME = Path(__file__).absolute().resolve().parent
PROJ_PATH = DIRNAME.parent.parent
# class T(SingletonBaseClass): TODO: singleton does not support args now.
class T:
"""Use the simplest way to (C)reate a Template and (r)ender it!!"""
def __init__(self, uri: str):
"""
here are some uri usages
case 1) "a.b.c:x.y.z"
It will load DIRNAME/a/b/c.yaml as `yaml` and load yaml[x][y][z]
case 2) ".c:x.y.z"
It will load c.yaml in caller's (who call `T(uri)`) directory as `yaml` and load yaml[x][y][z]
the loaded content will be saved in `self.template`
"""
# Inspect the calling stack to get the caller's directory
stack = inspect.stack()
caller_frame = stack[1]
caller_module = inspect.getmodule(caller_frame[0])
caller_dir = Path(caller_module.__file__).parent
# Parse the URI
path_part, yaml_path = uri.split(":")
yaml_keys = yaml_path.split(".")
if path_part.startswith("."):
yaml_file_path = caller_dir / f"{path_part[1:].replace('.', '/')}.yaml"
else:
yaml_file_path = (PROJ_PATH / path_part.replace(".", "/")).with_suffix(".yaml")
# Load the YAML file
with open(yaml_file_path, "r") as file:
yaml_content = yaml.safe_load(file)
# Traverse the YAML content to get the desired template
for key in yaml_keys:
yaml_content = yaml_content[key]
self.template = yaml_content
def r(self, **context: Any):
"""
Render the template with the given context.
"""
return Environment(undefined=StrictUndefined).from_string(self.template).render(**context)