Update agentic APIs

PiperOrigin-RevId: 814823652

Author: hgao327

tunix/rl/experimental/agentic/parser/tool_parser/gemini_parser.py

@@ -6,7 +6,7 @@
 from tunix.rl.experimental.agentic.tools import base_tool
 
 BaseTool = base_tool.BaseTool
-ToolCall = tool_parser_base.ToolCall
+ToolCall = base_tool.ToolCall
 ToolParser = tool_parser_base.ToolParser
 
 

tunix/rl/experimental/agentic/parser/tool_parser/qwen_parser.py

@@ -5,7 +5,7 @@
 """
 
 import json
-from typing import Any, List
+from typing import Any
 
 from tunix.rl.experimental.agentic.parser.tool_parser import tool_parser_base
 from tunix.rl.experimental.agentic.tools import base_tool
@@ -36,13 +36,10 @@ def parse(self, model_response: str) -> list[ToolCall]:
         model_response (str): Text containing tool calls
 
     Returns:
-        ToolInputs: Parsed tool calls
+        list[ToolCall]: Parsed tool calls
     """
     tool_calls_dicts = self.parse_qwen_tool_calls(model_response)
-    tool_calls = [
-        ToolCall(name=tc["name"], arguments=tc["arguments"])
-        for tc in tool_calls_dicts
-    ]
+    tool_calls = [ToolCall(**tool_call) for tool_call in tool_calls_dicts]
     return tool_calls
 
   def parse_qwen_tool_calls(self, text: str) -> list[dict[str, Any]]:
@@ -92,7 +89,7 @@ def parse_qwen_tool_calls(self, text: str) -> list[dict[str, Any]]:
 
   def get_tool_prompt(
       self,
-      tools: List[BaseTool],
+      tools: list[BaseTool],
       *,
       schema_style: str = "openai",
   ) -> str:

tunix/rl/experimental/agentic/parser/tool_parser/tool_parser_base.py

@@ -13,12 +13,6 @@
 ABC = abc.ABC
 
 
-@dataclass
-class ToolCall:
-  name: str
-  arguments: dict[str, Any]
-
-
 class ToolParser(ABC):
   """Abstract base class for all tool parsers.
 
@@ -28,7 +22,7 @@ class ToolParser(ABC):
   """
 
   @abstractmethod
-  def parse(self, model_response: str) -> list[ToolCall]:
+  def parse(self, model_response: str) -> list[base_tool.ToolCall]:
     """Parse model output and return a list of tool calls.
 
     Args:
@@ -42,17 +36,13 @@ def parse(self, model_response: str) -> list[ToolCall]:
   @abstractmethod
   def get_tool_prompt(
       self,
-      tools: List[BaseTool],
-      *,
-      schema_style: Literal["openai", "mcp", "gemini"] = "openai",
+      tools_schema: str,
   ) -> str:
     """Generate tool-usage instruction prompt from a list of tools.
 
     Args:
-        tools: List of tool instances (BaseTool).
-        schema_style: "openai" -> use tool.json (OpenAI function-calling style)
-          "mcp"    -> use tool.to_mcp_json() (MCP-compatible format) "gemini" ->
-          use Gemini-compatible schema
+        tools_schema: A string containing the tool schemas, generated by
+          `_tools_schema_dump`.
 
     Returns:
         str: Prompt text to feed into the model (includes tool schemas).
@@ -79,9 +69,9 @@ def _tools_schema_dump(
       schemas = [t.to_mcp_json() for t in tools]
     elif schema_style == "gemini":
       # Gemini also uses JSON schema, same as OpenAI
-      schemas = [t.json for t in tools]
+      schemas = [t.get_json_schema() for t in tools]
     else:
-      schemas = [t.json for t in tools]
+      schemas = [t.get_json_schema() for t in tools]
     return json.dumps(schemas, ensure_ascii=False, indent=2)
 
   def parse_tool_outputs(self) -> dict[str, Any]: