Code Refactoring

.github/workflows/cpu_tests.yml

@@ -9,18 +9,19 @@ on:
 jobs:
   test-envs:
     runs-on: ubuntu-latest
-    timeout-minutes: 15
+    timeout-minutes: 20
 
     strategy:
       matrix:
+        python-version: ["3.10"]
         test-file:
-          - tests/unit/envs/ --ignore tests/unit/envs/test_webshop_text_env.py --ignore tests/unit/envs/test_alfworld_env.py
-          - tests/unit/envs/test_alfworld_env.py
+          - agentfly/tests/unit/envs/ --ignore agentfly/tests/unit/envs/test_webshop_text_env.py --ignore agentfly/tests/unit/envs/test_alfworld_env.py
+          - agentfly/tests/unit/envs/test_alfworld_env.py
           # - tests/unit/envs/test_webshop_text_env.py # TODO: add minimal variant of the webshop docker image
-          - tests/unit/rewards/ --ignore tests/unit/rewards/test_env_id.py --ignore tests/unit/rewards/test_webshop_reward.py
-          - tests/unit/tools/ --ignore tests/unit/tools/test_webshop_tool.py --ignore tests/unit/tools/test_scienceworld_tool.py --ignore tests/unit/tools/test_code_tool.py
-          - tests/unit/tools/test_scienceworld_tool.py
-          - tests/unit/tools/test_code_tool.py
+          - agentfly/tests/unit/rewards/ --ignore agentfly/tests/unit/rewards/test_env_id.py --ignore agentfly/tests/unit/rewards/test_webshop_reward.py
+          - agentfly/tests/unit/tools/ --ignore agentfly/tests/unit/tools/test_webshop_tool.py --ignore agentfly/tests/unit/tools/test_scienceworld_tool.py --ignore agentfly/tests/unit/tools/test_code_tool.py
+          - agentfly/tests/unit/tools/test_scienceworld_tool.py
+          - agentfly/tests/unit/tools/test_code_tool.py
           # - test/unit/agents/ # TODO: recheck this
 
     steps:
@@ -34,6 +35,13 @@ jobs:
         with:
           python-version: '3.10'
 
+      - name: Verify Python
+        run: |
+          which python
+          python --version
+          which pip
+          python -m pip --version
+
       - name: Free up disk space
         run: |
           echo "Before cleanup:"
@@ -50,8 +58,9 @@ jobs:
 
       - name: Install dependencies (main repo)
         run: |
-          pip install -r agents/requirements.txt
+          pip install -e .
           pip install datasets
+          pip install -e '.[dev]' --no-build-isolation
 
       - name: Cache AgentFly cache
         uses: actions/cache@v4
@@ -75,5 +84,4 @@ jobs:
 
       - name: Run unit test (${{ matrix.test-file }})
         run: |
-          cd agents
-          python -m pytest ${{ matrix.test-file }}
+          pytest ${{ matrix.test-file }}

.gitignore

@@ -122,7 +122,7 @@ tests/e2e/toy_examples/deepspeed/synchronous/output.txt
 
 # data
 *.parquet
-agents/agents/data/*
+agentfly/agents/data/*
 
 # local logs
 logs
@@ -131,10 +131,12 @@ outputs
 *.out
 
 # Notebooks
-agents/tests/*.ipynb
-agents/tests/*.jpg
-agents/tests/*.jpeg
-agents/tests/*.png
-agents/agents/*.ipynb
-agents/temp/
+agentfly/tests/*.ipynb
+agentfly/tests/*.jpg
+agentfly/tests/*.jpeg
+agentfly/tests/*.png
+agentfly/agents/*.ipynb
+agentfly/temp/
+agentfly/data/
+*.ipynb
 

README.md

@@ -48,30 +48,28 @@ bash install.sh # Assume conda with python3.10.x
 ```
 **Option 2**: Customized Installation
 
-Clone and initialize the project:
-```bash
-git clone https://github.com/Agent-One-Lab/AgentFly
-cd AgentFly
-git submodule init
-git submodule update
-```
-Basic python packages installation:
-```bash
-pip install -e .
-pip install -e '.[verl]' --no-build-isolation
-```
-Optionally, some tools actually require some additional dependencies:
-
-Some of our tools & environments are managed by *enroot* backend. To use them, please install [enroot](https://github.com/NVIDIA/enroot/blob/master/doc/installation.md) (sudo required). Such tools include code_interpreter, retrieval, webshop, alfworld, sciencworld.
-
-Search requires redis to cache results, an optional way to install with conda:
-```bash
-conda install conda-forge::redis-server==7.4.0
-```
+Please refer to [installation.md](docs/start/installation.md) for custmoized installation.
 
 ## Quick Start
-```
+```python
+# Really small example to build an agent and run
+from agentfly.agents import HFAgent
+from agentfly.tools import calculate
+messages = [{"role": "user", "content": "What is the result of 1 + 1?"}]
+agent = HFAgent(
+    model_name_or_path="Qwen/Qwen2.5-3B-Instruct",
+    tools=[calculate],
+    template="qwen2.5",
+    backend="async_vllm",
+)
+await agent.run(
+    messages=messages,
+    max_turns=3,
+    num_chains=1
+)
 
+trajectories = agent.trajectories
+print(trajectories)
 ```
 
 ## Features

agents/agents/agents/__init__.py → agentfly/agents/__init__.py

@@ -2,5 +2,5 @@
 from .specialized.code_agent import CodeAgent
 from .specialized.think_agent import ThinkAgent
 from .specialized.gui_agent import GUIAgent
-
-__all__ = ["ReactAgent", "CodeAgent", "ThinkAgent", "GUIAgent"]
+from .specialized.hf_agent import HFAgent
+from .templates.utils import process_vision_info, tokenize_conversation, tokenize_conversations

agents/agents/agents/agent_base.py → agentfly/agents/agent_base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
 import json
-
+from .utils.messages import MessagesList
 from .templates.templates import get_template
 from ..__init__ import AGENT_DATA_DIR
 from .llm_backend import AsyncVLLMBackend, AsyncVerlBackend, ClientBackend, TransformersBackend, VLLMBackend
@@ -11,7 +11,7 @@
 import torch
 from .templates.utils import tokenize_conversations
 from .templates.vision_processor import is_vision_template
-from .chain.chain_base import ChainGeneration
+from .chain.chain_base import ChainRollout
 import os
 import transformers
 import warnings
@@ -27,7 +27,7 @@
 
 Logger = logging.getLogger(__name__)
 
-class BaseAgent(ChainGeneration, ABC):
+class BaseAgent(ChainRollout, ABC):
     """
     Base class for all agents. All agent should subclass this class. A customized agent can implement the following methods:
 
@@ -43,7 +43,7 @@ def __init__(
         system_prompt: str = None,
         tools: List = None,
         max_length: int=8192,
-        backend: str = "transformers",
+        backend: str = "async_vllm",
         backend_config: Any = None,
         reward_fn: Callable = None,
         log_file: str = "agent",
@@ -156,6 +156,47 @@ def _init_llm_engine(self, model_name_or_path: str, backend: str):
 
         return llm_engine
 
+    def _preprocess_messages(self, messages: List[Dict]):
+        """
+        Do some necessary preprocessings to the messages, such as adding the sytem prompt
+        Args:
+            messages: List of messages to preprocess.
+
+        Returns:
+            List of preprocessed messages.
+        """
+        messages_list = MessagesList.from_data(messages)
+        for messages in messages_list:
+            if self.system_prompt:
+                messages.set_system_prompt(self.system_prompt, enforce=False)
+
+        return messages_list.to_list()
+
+    async def run(self,
+        messages: Union[List[dict], np.ndarray, Dict],
+        max_turns: int,
+        generation_config: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        """
+        This is the main interface for running the agent. It is a wrapper of different 
+        rollout methods, which must be asynchronous. Currently, we only support chain-based rollout.
+        Args:
+            messages: List of messages to generate responses for.
+            max_turns: The maximum number of turns to generate.
+            generation_config: The generation configuration.
+            **kwargs: Additional keyword arguments for generation.
+
+        """
+        processed_messages = self._preprocess_messages(messages)
+
+        return await self.run_async(
+            processed_messages,
+            max_turns=max_turns,
+            generation_config=generation_config,
+            **kwargs,
+        )
+
     def set_llm_engine(self, llm_engine: Any, tokenizer: Any, processor: Any):
         assert self.backend == "async_verl", "Only async verl backend is supported for now"
 

agents/agents/agents/auto.py → agentfly/agents/auto.py

@@ -7,6 +7,7 @@
 from .react.react_agent import ReactAgent
 from .specialized.code_agent import CodeAgent
 from .specialized.gui_agent import GUIAgent
+from .specialized.hf_agent import HFAgent
 from ..rewards.reward_base import get_reward_from_name
 
 
@@ -167,4 +168,5 @@ def from_pretrained(
 AutoAgent.register_agent("code", CodeAgent)
 AutoAgent.register_agent("openai", OpenAIAgent)
 AutoAgent.register_agent("think", ThinkAgent)
-AutoAgent.register_agent("gui", GUIAgent)
+AutoAgent.register_agent("gui", GUIAgent)
+AutoAgent.register_agent("hf", HFAgent)