Enhance AgentJet framework and update documentation

- Updated README.md to clarify AgentJet's capabilities and added swarm training instructions.
- Refactored job.py to import ray only when needed.
- Fixed file opening in doc_reader.py, tracing_reader, and native_parallel_worker.py to include UTF-8 encoding.
- Improved as_oai_baseurl_apikey.py by using TYPE_CHECKING for conditional imports.
- Moved generate_auth_token function to interchange_utils.py and added API_KEY_PREFIX.
- Updated various file opening methods across the codebase to ensure UTF-8 encoding.
- Added new dependencies (hydra-core, datasets) to pyproject.toml.
- Modified deep_finance_judge.py to improve code readability and structure.
- Updated math.py to change the remote training model path for consistency.

Author: binary-husky

README.md

@@ -12,25 +12,31 @@
 </div>
 
 
-**AgentJet (AJet)** is a cutting-edge, user-friendly training framework designed to optimize agents and workflows (built with OpenAI SDK, AgentScope, Langchain, or just HTTP requests), fine-tuning language model weights behind the scenes.
+**AgentJet (AJet)** is a cutting-edge, user-friendly agent RL training framework designed to optimize agents and agentic workflows (supporting any agent built with OpenAI SDK, AgentScope, Langchain, or raw HTTP requests), fine-tuning LLM weights to enhance model performance.
 
-Simply provide your agent **workflow**, training **dataset**, and **reward** function, and **AgentJet** will be ready to enhance your agents to their optimal performance!
+**AgentJet (AJet)** has fully-distributed **swarm training** capability, which means that you can **deploy `ajet-swarm start` in GPU server(s) and then start training agents in your laptop(s)**! Simply provide your agent workflow, training dataset, and reward function, and AgentJet will be ready to go!
 
 
 
 ## ✈️ Minimum Example
 
-Let's begin with the simplest example: a math agent with a tool call.
+### Classic Mode
 
-- First, please check out the [installation guide](https://modelscope.github.io/AgentJet/en/installation/) to set up the training environment.
-- Then, tune your first model using the minimum example.
-  ```python
-  ajet --conf tutorial/example_math_agent/math_agent.yaml --backbone='verl'
+Let's begin with the simplest example: a math agent with a tool call. This is a simple & centralized training example.
 
-  # change to --backbone='trinity' if you want to switch to trinity training engine;
-  # or --backbone='debug' if you want to debug with only vLLM
-  ```
+1. please check out the [installation guide](https://modelscope.github.io/AgentJet/en/installation/) to set up the training environment.
+2. tune your first model using the minimum example.
+    ```python
+    ajet --conf ./tutorial/example_math_agent/math_agent.yaml --backbone='verl'
+    ```
+
+### Swarm Mode
 
+1. Start swarm server and begin swarm overwatch: `ajet-swarm start` and `ajet-swarm overwatch --swarm-url=http://localhost:10086`.
+2. From another device (or localhost), run [this script to train](https://github.com/modelscope/AgentJet/blob/main/tutorial/example_math_swarm/math.py):
+    ```python
+    AJET_SWARM_URL="http://swarm-server-ip:10086" python ./tutorial/example_math_swarm/math.py
+    ```
 
 ## ✈️ Features
 

ajet/copilot/job.py

@@ -12,7 +12,6 @@
 from types import SimpleNamespace
 from typing import Any, Callable, Union
 
-import ray
 import yaml
 from loguru import logger
 
@@ -138,6 +137,7 @@ def set_data(
         return self
 
     def tune(self, *args, **kwargs) -> "AgentJetJob":
+        import ray
         ast_cfg = self.config.ajet
         if not ast_cfg.rollout or not ast_cfg.rollout.user_workflow:
             raise ValueError("Workflow must be set via set_workflow before tuning.")

ajet/task_reader/document_reader/doc_reader.py

@@ -72,7 +72,7 @@ def _calculate_file_hash(self, file_path: str) -> str:
         """Calculate SHA256 hash of a file."""
         try:
             hash_sha256 = hashlib.sha256()
-            with open(file_path, "rb") as f:
+            with open(file_path, "rb", encoding="utf-8") as f:
                 for chunk in iter(lambda: f.read(4096), b""):
                     hash_sha256.update(chunk)
             return hash_sha256.hexdigest()

ajet/task_reader/tracing_reader/__init__.py

@@ -53,7 +53,7 @@ def _load_existing_tasks(self, path: str) -> List[Task]:
         if not os.path.exists(path):
             return []
         tasks: List[Task] = []
-        with open(path, "r") as f:
+        with open(path, "r", encoding="utf-8") as f:
             for line in f:
                 line = line.strip()
                 if not line:
@@ -66,7 +66,7 @@ def _append_tasks(self, path: str, tasks: List[Task]) -> None:
         if not tasks:
             return
         mode = "a" if os.path.exists(path) else "w"
-        with open(path, mode) as f:
+        with open(path, mode, encoding="utf-8") as f:
             for task in tasks:
                 obj = task.model_dump()
                 f.write(json.dumps(obj, ensure_ascii=False) + "\n")

ajet/task_rollout/native_parallel_worker.py

@@ -85,7 +85,7 @@ def _write_swarm_rollout_dynamic_log(self, observation_window):
         string_buffer = ""
         for info in observation_window["info"]:
             string_buffer += f"{info}\n"
-        with open(fp, "w") as f:
+        with open(fp, "w", encoding="utf-8") as f:
             f.write(string_buffer)
         return
 

ajet/tuner_lib/as_oai_baseurl_apikey.py

@@ -1,13 +1,12 @@
 import os
-from typing import Any
+from typing import Any, TYPE_CHECKING
 from pydantic import BaseModel, Field
-from ajet.context_tracker.multiagent_tracking import (
-    MultiAgentContextTracker,
-)
 from openai.resources.chat.chat import AsyncChat
 from openai.resources.completions import AsyncCompletions
-from .experimental.as_oai_model_client import generate_auth_token
+from ajet.tuner_lib.experimental.interchange_utils import generate_auth_token
 
+if TYPE_CHECKING:
+    from ajet.context_tracker.multiagent_tracking import MultiAgentContextTracker
 
 class MockAsyncCompletions(AsyncCompletions):
     async def create(self, *args, **kwargs) -> Any: # type: ignore
@@ -43,7 +42,7 @@ class OpenaiClientBaseUrlTuner(BaseModel):
     def __init__(
         self,
         config,
-        context_tracker: MultiAgentContextTracker,
+        context_tracker: "MultiAgentContextTracker",
         target_tag: str,
         agent_name: str,
         episode_uuid: str,

ajet/tuner_lib/experimental/as_oai_model_client.py

@@ -10,12 +10,14 @@
 
 from loguru import logger
 from typing import TYPE_CHECKING
-from vllm.entrypoints.openai.protocol import ChatCompletionRequest
 from openai.types.chat.chat_completion import ChatCompletion
-from ajet.tuner_lib.experimental.as_oai_model_server import InterchangeCompletionRequest, API_KEY_PREFIX
+from ajet.tuner_lib.experimental.as_oai_model_server import InterchangeCompletionRequest
 from ajet.utils.thread_executors import SharedInferenceTrackerThreadExecutor, SharedInterchangeThreadExecutor
 from ajet.tuner_lib.experimental.interchange_utils import get_zmq_socket
-from ajet.tuner_lib.experimental.interchange_utils import DEBUG
+from ajet.tuner_lib.experimental.interchange_utils import DEBUG, API_KEY_PREFIX
+
+if TYPE_CHECKING:
+    from vllm.entrypoints.openai.protocol import ChatCompletionRequest
 
 context = zmq.Context()
 atexit.register(context.term)
@@ -24,38 +26,6 @@
     from ajet.context_tracker.multiagent_tracking import MultiAgentContextTracker
 
 
-def generate_auth_token(agent_name, target_tag, episode_uuid, episode_address):
-    """
-    Generate a Base64-encoded auth_token from the given agent_name, target_tag, and episode_uuid.
-
-    Args:
-        agent_name (str): The name of the agent.
-        target_tag (str): The target tag.
-        episode_uuid (str): The UUID of the episode.
-
-    Returns:
-        str: The generated auth_token in the format "Bearer <base64_encoded_string>".
-    """
-    # Step 1: Construct the auth_data dictionary
-    auth_data = {
-        "agent_name": agent_name,
-        "target_tag": target_tag,
-        "episode_uuid": episode_uuid,
-        "episode_address": episode_address,
-    }
-
-    # Step 2: Convert the dictionary to a JSON string
-    json_string = json.dumps(auth_data)
-
-    # Step 3: Encode the JSON string into Base64
-    base64_encoded = base64.b64encode(json_string.encode('utf-8')).decode('utf-8')
-
-    # Step 4: Prepend "Bearer " to the Base64-encoded string
-    auth_token = f"{API_KEY_PREFIX}{base64_encoded}"    # API_KEY_PREFIX: Literal['sk-ajet-']
-
-    return auth_token
-
-
 class InterchangeClient:
     """ InterchangeClient is re-created in each episode
     """

ajet/tuner_lib/experimental/as_oai_model_server.py

@@ -35,9 +35,8 @@
 
 from ajet.utils.networking import get_host_ip
 from ajet.tuner_lib.experimental.interchange_utils import EpisodeStatus
-from ajet.tuner_lib.experimental.interchange_utils import DEBUG, VERBOSE
+from ajet.tuner_lib.experimental.interchange_utils import DEBUG, VERBOSE, API_KEY_PREFIX
 
-API_KEY_PREFIX = "sk-ajet-"
 
 class InterchangeCompletionRequest(BaseModel):
     completion_request: ChatCompletionRequest

ajet/tuner_lib/experimental/as_swarm_server.py

@@ -243,7 +243,7 @@ def _write_swarm_server_dynamic_log(shared_mem_dict):
                     p = es.model_dump_json()
                     string_buffer += f"{p}\n"
 
-            with open(fp, "w") as f:
+            with open(fp, "w", encoding="utf-8") as f:
                 f.write(string_buffer)
         return
 

ajet/tuner_lib/experimental/interchange_utils.py

@@ -1,6 +1,9 @@
 import os
 import time
 import httpx
+import base64
+import json
+
 from typing import List
 from pydantic import BaseModel, Field
 from loguru import logger
@@ -19,6 +22,7 @@
     "ENGINE.WEIGHT_EXPORTING"
 ]
 
+API_KEY_PREFIX = "sk-ajet-"
 
 class SyncTrainConfigRequest(BaseModel):
     yaml_as_string: str
@@ -205,3 +209,36 @@ def get_zmq_socket(config, episode_uuid: str, tag: str = ""):
     else:
         raise RuntimeError(f"Unknown interchange_method: {interchange_method}")
     return zmq_contect_address, ipc_path
+
+
+
+def generate_auth_token(agent_name, target_tag, episode_uuid, episode_address):
+    """
+    Generate a Base64-encoded auth_token from the given agent_name, target_tag, and episode_uuid.
+
+    Args:
+        agent_name (str): The name of the agent.
+        target_tag (str): The target tag.
+        episode_uuid (str): The UUID of the episode.
+
+    Returns:
+        str: The generated auth_token in the format "Bearer <base64_encoded_string>".
+    """
+    # Step 1: Construct the auth_data dictionary
+    auth_data = {
+        "agent_name": agent_name,
+        "target_tag": target_tag,
+        "episode_uuid": episode_uuid,
+        "episode_address": episode_address,
+    }
+
+    # Step 2: Convert the dictionary to a JSON string
+    json_string = json.dumps(auth_data)
+
+    # Step 3: Encode the JSON string into Base64
+    base64_encoded = base64.b64encode(json_string.encode('utf-8')).decode('utf-8')
+
+    # Step 4: Prepend "Bearer " to the Base64-encoded string
+    auth_token = f"{API_KEY_PREFIX}{base64_encoded}"    # API_KEY_PREFIX: Literal['sk-ajet-']
+
+    return auth_token