refactor: clean up unused imports and add text extraction method in MultiAgentContextTracker; add skip GPU check option in launcher

Author: binary-husky

ajet/context_tracker/multiagent_tracking.py

@@ -20,15 +20,7 @@
 from ajet.utils.color_hsl import adjust_color_hsl
 from ajet.utils.compute_madness import compute_string_madness
 from ajet.utils.tokenizer import ajet_apply_chat_template
-#  
-from ajet.utils.msg_converter import (
-    convert_grouped_steps_to_openai_format,
-    convert_ext_msg_to_openai_format,
-    agentscope_to_openai,
-    openai_to_agentscope,
-    agentscope_to_openai_grouped,
-    openai_to_agentscope_grouped,
-)
+
 @dataclass
 class TimelineMergingPolicyConfig:
     timeline_compare_level: str = "text"
@@ -82,6 +74,40 @@ def preprocess_tools_field(self, tools: List[dict] = [], disable_toolcalls: bool
                     tools[i]["function"]["parameters"] = tools[i]["function"].pop("parameters")
         return tools
 
+    def extract_text_content_from_content_dict(self, msg):
+        # msg = {
+        #    "role": "assistant",
+        #    "content": [
+        #        {
+        #           "type": "text",
+        #           "text": "some text"
+        #        },
+        #    ],
+        # }
+
+        str_content = ""
+        for item in msg["content"]:
+            # item = {
+            #   "type": "text",
+            #   "text": "some text"
+            # },
+
+            assert isinstance(item, dict), f"Unsupported non-dict item in message content: {item}. Full message: {msg}"
+
+            if ("text" not in item):
+                logger.warning(
+                    f"Non-text content in message content detected: {item}. Ignoring."
+                )
+                should_skip_message = True
+                return str_content, should_skip_message
+
+            if isinstance(item["text"], str):
+                str_content += str(item["text"])
+            else:
+                str_content = ""
+
+        should_skip_message = False
+        return str_content, should_skip_message
 
     def step_spawn_timeline(self, messages: List[dict], tools: List = [], disable_toolcalls: bool = False) -> List[ExtendedMessage]:
         """Spawn a timeline from messages.
@@ -101,55 +127,32 @@ def step_spawn_timeline(self, messages: List[dict], tools: List = [], disable_to
             consider_roles.remove("tool")
 
         for i, msg in enumerate(messages):
+
             if (disable_toolcalls) and (not isinstance(msg["content"], str)):
                 continue
+
             if msg["role"] not in consider_roles:
                 continue
+
             if not isinstance(msg["content"], str):
                 author = "env"
-                ignore = False
-                str_content = ""
-                extracted_tool_call_id = ""
-                for item_idx, item in enumerate(msg["content"]):
-                    if isinstance(item, dict) and item.get("type") == "tool_result":
-                        is_tool_result_msg = True  # 标记为 tool_result 消息
-                        # Extract tool_call_id from the tool_result block
-                        if item.get("id"):
-                            extracted_tool_call_id = item.get("id", "")
-                        output = item.get("output", "")
-                        if isinstance(output, str):
-                            str_content += output
-                        elif isinstance(output, list):
-                            # output can be List[TextBlock | ImageBlock | AudioBlock]
-                            for out_item in output:
-                                if isinstance(out_item, str):
-                                    str_content += out_item
-                                elif isinstance(out_item, dict) and "text" in out_item:
-                                    str_content += str(out_item["text"])
-                        else:
-                            str_content += str(output)
-                    elif isinstance(item, dict) and "text" in item:
-                        if isinstance(item["text"], str):
-                            str_content += str(item["text"])
-                        else:
-                            str_content = ""
-                    else:
-                        logger.warning(
-                            f"Non-text content in message content detected: {item}. Ignoring."
-                        )
-                        ignore = True
-                        break
-                msg["content"] = str_content
-                msg["tool_call_id"] = extracted_tool_call_id  # Store extracted tool_call_id
-                
-                # ★ 关键修复：如果是 tool_result 消息，将 role 恢复为 "tool"（OpenAI 格式）
-                if is_tool_result_msg and extracted_tool_call_id:
-                    msg["role"] = "tool"
-                
-
-                if ignore:
+                should_skip_message = False
+
+                # fix msg content
+                if msg["content"] is None:
+                    msg["content"] = ""
+
+                elif isinstance(msg["content"], list):
+                    msg["content"], should_skip_message = self.extract_text_content_from_content_dict(msg)
+
+                else:
+                    raise ValueError(f"Unsupported non-str message content type: {type(msg['content'])}, Message:\n {msg}")
+
+                if should_skip_message:
                     continue
-                msg["content"] = str(msg["content"])  # TODO: better handling mm data
+
+                if not isinstance(msg["content"], str):
+                    msg["content"] = str(msg["content"])  # TODO: better handling mm data
 
             if msg["role"] == "system":
                 author = "initialization"
@@ -169,6 +172,7 @@ def step_spawn_timeline(self, messages: List[dict], tools: List = [], disable_to
                     tool_calls=(msg["tool_calls"] if "tool_calls" in msg else []),
                     tool_call_id=(msg["tool_call_id"] if "tool_call_id" in msg else ""),
                     token_generator="auto",
+                    name = (msg["name"] if "name" in msg else ""),
                     first_message=(i == 0),
                 )
             ]
@@ -605,25 +609,3 @@ def check_context_token_num_safe(
         else:
             ret = (False, token_overflow, "token_overflow")
         return ret
-
-    def get_grouped_steps_openai_format(self) -> List[List[Dict[str, Any]]]:
-        """
-        将 grouped_steps 转换为 OpenAI 格式并返回。
-        
-        Returns:
-            OpenAI 格式的轨迹数据 (List of List of dict)
-            每条消息格式如：
-            - {"role": "assistant", "content": "...", "tool_calls": [...]}
-            - {"role": "tool", "content": "...", "tool_call_id": "call_xxx"}
-            - {"role": "user/system", "content": "..."}
-        """
-        return convert_grouped_steps_to_openai_format(self.grouped_steps)
-
-    def get_full_context_openai_format(self) -> List[Dict[str, Any]]:
-        """
-        将当前 full_context 转换为 OpenAI 格式并返回。
-        
-        Returns:
-            OpenAI 格式的消息列表 (List of dict)
-        """
-        return [convert_ext_msg_to_openai_format(msg) for msg in self.full_context]

ajet/launcher.py

@@ -86,6 +86,7 @@ def parse_args():
         help="Launch Crafters Env Simulation",
     )
     parser.add_argument("--reboot", action="store_true", default=False, help="reboot flag")
+    parser.add_argument("--skip-check-avail-gpu", action="store_true", default=False, help="Skip GPU availability check")
     parser.add_argument(
         "--kill",
         type=str,
@@ -305,7 +306,7 @@ def main():
 
     if args.with_finworld:
         pty_launch("finworld")
-        
+
     if args.with_crafters:
         pty_launch("crafters")
 

ajet/schema/extended_msg.py

@@ -74,6 +74,7 @@ def __init__(
         tool_calls=[],
         tool_call_id="",
         token_logprob_arr=[],
+        name="",    # preserved field, not used currently
         first_message=False,
     ):
         self.author = author
@@ -90,6 +91,7 @@ def __init__(
         self.tools = tools
         self.tool_calls = tool_calls
         self.tool_call_id = tool_call_id
+        self.name = name    # preserved field, not used currently
         if not isinstance(self.tool_calls, list):
             # agent scope sometimes gives weird type for tool_calls, which is against OpenAI schema
             self.tool_calls = list(self.tool_calls)