diff --git a/rdagent/components/proposal/prompts.yaml b/rdagent/components/proposal/prompts.yaml index 12de3cd0..4e3b427d 100644 --- a/rdagent/components/proposal/prompts.yaml +++ b/rdagent/components/proposal/prompts.yaml @@ -39,4 +39,4 @@ hypothesis2experiment: {{ target_list }} To help you generate new {{targets}}, we have prepared the following information for you: {{ RAG }} - Please generate the new {{targets}} based on the information above. \ No newline at end of file + Please generate the new {{targets}} based on the information above. diff --git a/rdagent/utils/agent/__init__.py b/rdagent/utils/agent/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/rdagent/utils/agent/ret.py b/rdagent/utils/agent/ret.py new file mode 100644 index 00000000..8c8c29f4 --- /dev/null +++ b/rdagent/utils/agent/ret.py @@ -0,0 +1,33 @@ +""" +The output of a agent is very important. + +We think this part can be shared. +""" +from abc import abstractclassmethod +import re +from typing import Any + +from rdagent.utils.agent.tpl import T + + +class AgentOut: + @abstractclassmethod + def get_spec(cls, **context: Any) -> str: + raise NotImplementedError(f"Please implement the `get_spec` method") + + @classmethod + def extract_output(cls, resp: str) -> Any: + raise resp + + +class PythonAgentOut(AgentOut): + @classmethod + def get_spec(cls): + return T(".tpl:PythonAgentOut").r() + + @classmethod + def extract_output(cls, resp: str): + match = re.search(r".*```[Pp]ython\n(.*)\n```.*", resp, re.DOTALL) + if match: + code = match.group(1) + return code diff --git a/rdagent/utils/agent/tpl.py b/rdagent/utils/agent/tpl.py new file mode 100644 index 00000000..e7ec2c56 --- /dev/null +++ b/rdagent/utils/agent/tpl.py @@ -0,0 +1,63 @@ +""" +Here are some infrastruture to build a agent + +The motivation of tempalte and AgentOutput Design +""" + +from typing import Any +from jinja2 import Environment, StrictUndefined + +from pathlib import Path + +import yaml +import inspect + +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) diff --git a/rdagent/utils/agent/tpl.yaml b/rdagent/utils/agent/tpl.yaml new file mode 100644 index 00000000..2b41013a --- /dev/null +++ b/rdagent/utils/agent/tpl.yaml @@ -0,0 +1,6 @@ +PythonAgentOut: |- + The return code should be like + ```Python + + ``` + diff --git a/test/utils/test_agent_infra.py b/test/utils/test_agent_infra.py new file mode 100644 index 00000000..ea1ce25d --- /dev/null +++ b/test/utils/test_agent_infra.py @@ -0,0 +1,33 @@ +import unittest + +from rdagent.oai.llm_utils import APIBackend +from rdagent.utils.agent.ret import PythonAgentOut +from rdagent.utils.agent.tpl import T + + + + +class TestAgentInfra(unittest.TestCase): + + def test_agent_infra(self): + + # NOTE: It is not serious. It is just for testing + sys_prompt = T("components.proposal.prompts:hypothesis_gen.system_prompt").r( + targets="targets", + scenario=T("scenarios.qlib.experiment.prompts:qlib_model_background").r(), + hypothesis_output_format=PythonAgentOut.get_spec(), + hypothesis_specification=PythonAgentOut.get_spec(), + ) + user_prompt = T("components.proposal.prompts:hypothesis_gen.user_prompt").r( + hypothesis_and_feedback="No Feedback", + RAG="No RAG", + targets="targets", + ) + resp = APIBackend().build_messages_and_create_chat_completion(user_prompt=user_prompt, system_prompt=sys_prompt) + code = PythonAgentOut.extract_output(resp) + + print(code) + + +if __name__ == "__main__": + unittest.main()