feat(agent): support plain function as client_tool

Summary:

Test Plan:

Author: ehhuang

src/llama_stack_client/lib/agents/agent.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 import logging
-from typing import AsyncIterator, Iterator, List, Optional, Tuple, Union
+from typing import AsyncIterator, Callable, Iterator, List, Optional, Tuple, Union, Any
 
 from llama_stack_client import LlamaStackClient
 
@@ -18,7 +18,7 @@
 from llama_stack_client.types.shared_params.response_format import ResponseFormat
 from llama_stack_client.types.shared_params.sampling_params import SamplingParams
 
-from .client_tool import ClientTool
+from .client_tool import ClientTool, client_tool
 from .tool_parser import ToolParser
 
 DEFAULT_MAX_ITER = 10
@@ -28,10 +28,12 @@
 
 class AgentUtils:
     @staticmethod
-    def get_client_tools(tools: Optional[List[Union[Toolgroup, ClientTool]]]) -> List[ClientTool]:
+    def get_client_tools(tools: Optional[List[Union[Toolgroup, ClientTool, Callable[..., Any]]]]) -> List[ClientTool]:
         if not tools:
             return []
 
+        # Wrap any function in client_tool decorator
+        tools = [client_tool(tool) if (callable(tool) and not isinstance(tool, ClientTool)) else tool for tool in tools]
         return [tool for tool in tools if isinstance(tool, ClientTool)]
 
     @staticmethod
@@ -59,7 +61,8 @@ def get_turn_id(chunk: AgentTurnResponseStreamChunk) -> Optional[str]:
     def get_agent_config(
         model: Optional[str] = None,
         instructions: Optional[str] = None,
-        tools: Optional[List[Union[Toolgroup, ClientTool]]] = None,
+        tools: Optional[List[Union[Toolgroup, ClientTool, Callable[..., Any]]]] = None,
+        client_tools: Optional[List[ClientTool]] = None,
         tool_config: Optional[ToolConfig] = None,
         sampling_params: Optional[SamplingParams] = None,
         max_infer_iters: Optional[int] = None,
@@ -96,15 +99,12 @@ def get_agent_config(
             agent_config["tool_config"] = tool_config
         if tools is not None:
             toolgroups: List[Toolgroup] = []
-            client_tools: List[ClientTool] = []
-
             for tool in tools:
                 if isinstance(tool, str) or isinstance(tool, dict):
                     toolgroups.append(tool)
-                else:
-                    client_tools.append(tool)
 
             agent_config["toolgroups"] = toolgroups
+        if client_tools:
             agent_config["client_tools"] = [tool.get_tool_definition() for tool in client_tools]
 
         agent_config = AgentConfig(**agent_config)
@@ -122,7 +122,7 @@ def __init__(
         tool_parser: Optional[ToolParser] = None,
         model: Optional[str] = None,
         instructions: Optional[str] = None,
-        tools: Optional[List[Union[Toolgroup, ClientTool]]] = None,
+        tools: Optional[List[Union[Toolgroup, ClientTool, Callable[..., Any]]]] = None,
         tool_config: Optional[ToolConfig] = None,
         sampling_params: Optional[SamplingParams] = None,
         max_infer_iters: Optional[int] = None,
@@ -143,7 +143,7 @@ def __init__(
         :param instructions: The instructions for the agent.
         :param tools: A list of tools for the agent. Values can be one of the following:
             - dict representing a toolgroup/tool with arguments: e.g. {"name": "builtin::rag/knowledge_search", "args": {"vector_db_ids": [123]}}
-            - a python function decorated with @client_tool
+            - a python function with a docstring. See @client_tool for more details.
             - str representing a tool within a toolgroup: e.g. "builtin::rag/knowledge_search"
             - str representing a toolgroup_id: e.g. "builtin::rag", "builtin::code_interpreter", where all tools in the toolgroup will be added to the agent
             - an instance of ClientTool: A client tool object.
@@ -164,10 +164,12 @@ def __init__(
 
         # Construct agent_config from parameters if not provided
         if agent_config is None:
+            client_tools = AgentUtils.get_client_tools(tools)
             agent_config = AgentUtils.get_agent_config(
                 model=model,
                 instructions=instructions,
                 tools=tools,
+                client_tools=client_tools,
                 tool_config=tool_config,
                 sampling_params=sampling_params,
                 max_infer_iters=max_infer_iters,
@@ -176,7 +178,6 @@ def __init__(
                 response_format=response_format,
                 enable_session_persistence=enable_session_persistence,
             )
-            client_tools = AgentUtils.get_client_tools(tools)
 
         self.agent_config = agent_config
         self.client_tools = {t.get_name(): t for t in client_tools}
@@ -332,7 +333,7 @@ def __init__(
         tool_parser: Optional[ToolParser] = None,
         model: Optional[str] = None,
         instructions: Optional[str] = None,
-        tools: Optional[List[Union[Toolgroup, ClientTool]]] = None,
+        tools: Optional[List[Union[Toolgroup, ClientTool, Callable[..., Any]]]] = None,
         tool_config: Optional[ToolConfig] = None,
         sampling_params: Optional[SamplingParams] = None,
         max_infer_iters: Optional[int] = None,
@@ -353,7 +354,7 @@ def __init__(
         :param instructions: The instructions for the agent.
         :param tools: A list of tools for the agent. Values can be one of the following:
             - dict representing a toolgroup/tool with arguments: e.g. {"name": "builtin::rag/knowledge_search", "args": {"vector_db_ids": [123]}}
-            - a python function decorated with @client_tool
+            - a python function with a docstring. See @client_tool for more details.
             - str representing a tool within a toolgroup: e.g. "builtin::rag/knowledge_search"
             - str representing a toolgroup_id: e.g. "builtin::rag", "builtin::code_interpreter", where all tools in the toolgroup will be added to the agent
             - an instance of ClientTool: A client tool object.
@@ -374,10 +375,12 @@ def __init__(
 
         # Construct agent_config from parameters if not provided
         if agent_config is None:
+            client_toolss = AgentUtils.get_client_tools(tools)
             agent_config = AgentUtils.get_agent_config(
                 model=model,
                 instructions=instructions,
                 tools=tools,
+                client_tools=client_toolss,
                 tool_config=tool_config,
                 sampling_params=sampling_params,
                 max_infer_iters=max_infer_iters,
@@ -386,7 +389,6 @@ def __init__(
                 response_format=response_format,
                 enable_session_persistence=enable_session_persistence,
             )
-            client_tools = AgentUtils.get_client_tools(tools)
 
         self.agent_config = agent_config
         self.client_tools = {t.get_name(): t for t in client_tools}