add interchange api

Author: binary-husky

ajet/backbone/main_trinity.py

@@ -40,6 +40,8 @@ def patched_trainer_get_actor(cls, config: Config):
     Trainer.get_actor = classmethod(patched_trainer_get_actor)
 
 
+
+
 if __name__ == "__main__":
     patch_runtime_env_to_get_actor()
     main()

ajet/backbone/main_verl.py

@@ -247,6 +247,10 @@ def run(self, config):
 
         from ajet.backbone.trainer_verl import AjetRayPPOTrainer
 
+        if config.ajet.enable_experimental_reverse_proxy:
+            from ajet.tuner_lib.weight_tuner.experimental.as_oai_model_server import start_interchange_server
+            start_interchange_server()
+
         # Initialize the PPO trainer.
         trainer = AjetRayPPOTrainer(
             config=config,

ajet/backbone/trainer_trinity.py

@@ -19,7 +19,7 @@
 from trinity.utils.monitor import MONITOR, Monitor
 
 from ajet.backbone.warm_up import warm_up_process
-from ajet.context_tracker.agentscope_tracker.multiagent_tracking import (
+from ajet.context_tracker.multiagent_tracking import (
     MultiAgentContextTracker,
 )
 from ajet.schema.trajectory import Sample
@@ -116,6 +116,10 @@ def __init__(
 
     async def run_async(self):
         ajet_config = get_ajet_config_from_trinity_side()
+        if ajet_config.ajet.enable_experimental_reverse_proxy:
+            raise NotImplementedError(
+                "The experimental reverse proxy is not supported in Trinity backbone yet."
+            )
         warm_up_process(ajet_config)
         tracker = await TrinityRolloutManager(
             is_eval=self.is_eval,

…gentscope_tracker/multiagent_tracking.py …t/context_tracker/multiagent_tracking.pyajet/context_tracker/agentscope_tracker/multiagent_tracking.py renamed to ajet/context_tracker/multiagent_tracking.py ajet/context_tracker/agentscope_tracker/multiagent_tracking.py renamed to ajet/context_tracker/multiagent_tracking.py

@@ -8,7 +8,7 @@
 from loguru import logger
 from transformers.tokenization_utils import PreTrainedTokenizer
 
-from ajet.context_tracker.agentscope_tracker.timeline_merging import (
+from ajet.context_tracker.timeline_merging.timeline_merging import (
     merge_tracker_timelines, is_timeline_mergeable
 )
 from ajet.context_tracker.basic_tracker import (
@@ -36,6 +36,8 @@ class ContextTrackerConfig:
     detect_timeline_snap: bool = False
 
 
+
+
 class MultiAgentContextTracker(BaseContextTracker):
     """
     Context tracker is responsible to monitor and process LLM IO.
@@ -44,22 +46,22 @@ class MultiAgentContextTracker(BaseContextTracker):
 
     def __init__(
         self,
-        llm_inference_fn,
         tokenizer: PreTrainedTokenizer,
         config,
         should_interrupt_fn,
         generated_token_callback_fn,
+        episode_uuid: str,
         **kwargs,
     ):
         super().__init__(config, tokenizer, **kwargs)
-        self.llm_inference_fn = llm_inference_fn
         self.tokenizer = tokenizer
         self.should_interrupt_fn = should_interrupt_fn
         self.generated_token_callback_fn = generated_token_callback_fn
         self.context_overflow = False
         self.output_kwargs = {}
         self.input_kwargs = {}
         self.timeline_cache = {}
+        self.episode_uuid = episode_uuid
 
 
     def step_prepare(self, messages: List[dict], tools: List = [], timeline_uuid: str = ""):

…r/agentscope_tracker/timeline_merging.py …ker/timeline_merging/timeline_merging.pyajet/context_tracker/agentscope_tracker/timeline_merging.py renamed to ajet/context_tracker/timeline_merging/timeline_merging.py ajet/context_tracker/agentscope_tracker/timeline_merging.py renamed to ajet/context_tracker/timeline_merging/timeline_merging.py

@@ -6,6 +6,9 @@ ajet:
   backbone: debug # `debug` or `trinity` or `verl`
 
 
+  # the experimental reverse proxy feature that allows `tuner.as_oai_baseurl_apikey` feature
+  enable_experimental_reverse_proxy: True
+
   model:
     # which model should be trained
     path: /path/to/model/such/as/Qwen/Qwen2___5-14B-Instruct

ajet/default_config/ajet_default.yaml

@@ -1,8 +1,14 @@
 
+import time
 from openai.types.chat.chat_completion import ChatCompletion, Choice
 from openai.types.chat.chat_completion_message import ChatCompletionMessage
+from agentscope.model import ChatResponse as AgentScopeChatResponse
 from openai.types.completion_usage import CompletionUsage
-import time
+from typing import Any, Callable, Dict, List, Literal, Type, Union
+from agentscope.message import TextBlock, ToolUseBlock
+from agentscope._utils._common import _json_loads_with_repair
+from pydantic import BaseModel
+from agentscope.model import ChatResponse
 
 
 def convert_llm_proxy_response_to_oai_response(llm_proxy_response):
@@ -40,6 +46,66 @@ def convert_llm_proxy_response_to_oai_response(llm_proxy_response):
         usage=usage,
     )
 
+# copied from AgentScope's DashScopeChatModule
+def convert_llm_proxy_response_to_agentscope_response(
+    message,
+    structured_model: Type[BaseModel] | None = None,
+) -> AgentScopeChatResponse:    # type: ignore
+    content_blocks: List[TextBlock | ToolUseBlock] = []
+    content = message.get("content")
+    metadata: dict | None = None
+
+    if content not in [
+        None,
+        "",
+        [],
+    ]:
+        if isinstance(content, list):
+            for item in content:
+                if isinstance(item, dict) and "text" in item:
+                    content_blocks.append(
+                        TextBlock(
+                            type="text",
+                            text=item["text"],
+                        ),
+                    )
+        else:
+            content_blocks.append(
+                TextBlock(
+                    type="text",
+                    text=content,
+                ),
+            )
+
+    if message.get("tool_calls"):
+        for tool_call in message["tool_calls"]:
+            input_ = _json_loads_with_repair(
+                tool_call["function"].get(
+                    "arguments",
+                    "{}",
+                )
+                or "{}",
+            )
+            content_blocks.append(
+                ToolUseBlock(
+                    type="tool_use",
+                    name=tool_call["function"]["name"],
+                    input=input_,  # type: ignore
+                    id=tool_call["id"],
+                ),
+            )
+
+            if structured_model:
+                metadata = input_  # type: ignore
+
+    parsed_response = AgentScopeChatResponse(
+        content=content_blocks,
+        metadata=metadata,
+    )
+
+    return parsed_response
+
+
 
 def test_convert_llm_proxy_response_to_oai_response():
     """Test the conversion from llm_proxy_response to OpenAI ChatCompletion format."""

ajet/schema/convertion.py

@@ -26,7 +26,7 @@ class WorkflowTask(BaseModel):
     task_thread_index: int = Field(default=0)
     task_batch_index: int = Field(default=0)
     task_tag: str = Field(default="")
-    task_env_uuid: str = Field(default="")
+    episode_uuid: str = Field(default="")
     observation_window: dict = Field(default={})
     llm_inference_fn: Any = Field(default=None)
     tokenizer: Any = Field(default=None)

ajet/schema/task.py

@@ -10,7 +10,7 @@ def compute_reward(self, workflow_task: WorkflowTask, workflow_output: WorkflowO
         raw_reward = 0
 
         env = workflow_task.gym_env
-        raw_reward = env.evaluate(workflow_task.task_env_uuid, params={"sparse": False})
+        raw_reward = env.evaluate(workflow_task.episode_uuid, params={"sparse": False})
         if raw_reward >= 1:
             is_success = True
         else: