Agent Infra (#107)

* Agent Infra

* Update test/utils/test_agent_infra.py

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.
+    Please generate the new {{targets}} based on the information above.

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

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)

rdagent/utils/agent/tpl.yaml

@@ -0,0 +1,6 @@
+PythonAgentOut: |-
+  The return code should be like
+  ```Python
+  <You code>
+  ```
+  

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()