refactor, comments

Author: yanxi0830

src/llama_stack_client/lib/agents/react/agent.py

@@ -43,29 +43,36 @@ def __init__(
         json_response_format: bool = False,
         custom_agent_config: Optional[AgentConfig] = None,
     ):
-        def get_tool_definition(tool):
-            return {
-                "name": tool.identifier,
-                "description": tool.description,
-                "parameters": tool.parameters,
-            }
-
-        if custom_agent_config is None:
-            tool_names = ""
-            tool_descriptions = ""
+        def get_tool_defs():
+            tool_defs = []
             for x in builtin_toolgroups:
-                tool_names += ", ".join([tool.identifier for tool in client.tools.list(toolgroup_id=x)])
-                tool_descriptions += "\n".join(
-                    [f"- {tool.identifier}: {get_tool_definition(tool)}" for tool in client.tools.list(toolgroup_id=x)]
+                tool_defs.extend(
+                    [
+                        {
+                            "name": tool.identifier,
+                            "description": tool.description,
+                            "parameters": tool.parameters,
+                        }
+                        for tool in client.tools.list(toolgroup_id=x)
+                    ]
                 )
-
-            tool_names += ", "
-            tool_descriptions += "\n"
-            tool_names += ", ".join([tool.get_name() for tool in client_tools])
-            tool_descriptions += "\n".join(
-                [f"- {tool.get_name()}: {tool.get_tool_definition()}" for tool in client_tools]
+            tool_defs.extend(
+                [
+                    {
+                        "name": tool.get_name(),
+                        "description": tool.get_description(),
+                        "parameters": tool.get_params_definition(),
+                    }
+                    for tool in client_tools
+                ]
             )
+            return tool_defs
 
+        if custom_agent_config is None:
+            tool_names, tool_descriptions = "", ""
+            tool_defs = get_tool_defs()
+            tool_names = ", ".join([x["name"] for x in tool_defs])
+            tool_descriptions = "\n".join([f"- {x['name']}: {x}" for x in tool_defs])
             instruction = DEFAULT_REACT_AGENT_SYSTEM_PROMPT_TEMPLATE.replace("<<tool_names>>", tool_names).replace(
                 "<<tool_descriptions>>", tool_descriptions
             )

src/llama_stack_client/lib/agents/react/output_parser.py

@@ -4,13 +4,12 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError
 from typing import Dict, Any, Optional
 from ..output_parser import OutputParser
 from llama_stack_client.types.shared.completion_message import CompletionMessage
 from llama_stack_client.types.shared.tool_call import ToolCall
 
-import json
 import uuid
 
 
@@ -29,17 +28,17 @@ class ReActOutputParser(OutputParser):
     def parse(self, output_message: CompletionMessage) -> CompletionMessage:
         response_text = str(output_message.content)
         try:
-            response_json = json.loads(response_text)
-        except json.JSONDecodeError as e:
+            react_output = ReActOutput.model_validate_json(response_text)
+        except ValidationError as e:
             print(f"Error parsing action: {e}")
             return output_message
 
-        if response_json.get("answer", None):
+        if react_output.answer:
             return output_message
 
-        if response_json.get("action", None):
-            tool_name = response_json["action"].get("tool_name", None)
-            tool_params = response_json["action"].get("tool_params", None)
+        if react_output.action:
+            tool_name = react_output.action.tool_name
+            tool_params = react_output.action.tool_params
             if tool_name and tool_params:
                 call_id = str(uuid.uuid4())
                 output_message.tool_calls = [ToolCall(call_id=call_id, tool_name=tool_name, arguments=tool_params)]