fix(manual): update lib/ code for the input_schema changes

Author: ashwinb

src/llama_stack_client/lib/agents/client_tool.py

@@ -19,9 +19,17 @@
     Union,
 )
 
+from typing_extensions import TypedDict
+
 from llama_stack_client.types import CompletionMessage, Message
 from llama_stack_client.types.alpha import ToolResponse
-from llama_stack_client.types.tool_def_param import Parameter, ToolDefParam
+from llama_stack_client.types.tool_def_param import ToolDefParam
+
+
+class JSONSchema(TypedDict, total=False):
+    type: str
+    properties: Dict[str, Any]
+    required: List[str]
 
 
 class ClientTool:
@@ -47,28 +55,18 @@ def get_description(self) -> str:
         raise NotImplementedError
 
     @abstractmethod
-    def get_params_definition(self) -> Dict[str, Parameter]:
+    def get_input_schema(self) -> JSONSchema:
         raise NotImplementedError
 
     def get_instruction_string(self) -> str:
         return f"Use the function '{self.get_name()}' to: {self.get_description()}"
 
-    def parameters_for_system_prompt(self) -> str:
-        return json.dumps(
-            {
-                "name": self.get_name(),
-                "description": self.get_description(),
-                "parameters": {name: definition for name, definition in self.get_params_definition().items()},
-            }
-        )
-
     def get_tool_definition(self) -> ToolDefParam:
         return ToolDefParam(
             name=self.get_name(),
             description=self.get_description(),
-            parameters=list(self.get_params_definition().values()),
+            input_schema=self.get_input_schema(),
             metadata={},
-            tool_prompt_format="python_list",
         )
 
     def run(
@@ -148,6 +146,37 @@ def async_run_impl(self, **kwargs):
 T = TypeVar("T", bound=Callable)
 
 
+def _python_type_to_json_schema_type(type_hint: Any) -> str:
+    """Convert Python type hints to JSON Schema type strings."""
+    # Handle Union types (e.g., Optional[str])
+    origin = get_origin(type_hint)
+    if origin is Union:
+        # Get non-None types from Union
+        args = [arg for arg in get_args(type_hint) if arg is not type(None)]
+        if args:
+            type_hint = args[0]  # Use first non-None type
+
+    # Get the actual type if it's a generic
+    if hasattr(type_hint, "__origin__"):
+        type_hint = type_hint.__origin__
+
+    # Map Python types to JSON Schema types
+    type_name = getattr(type_hint, "__name__", str(type_hint))
+
+    type_mapping = {
+        "bool": "boolean",
+        "int": "integer",
+        "float": "number",
+        "str": "string",
+        "list": "array",
+        "dict": "object",
+        "List": "array",
+        "Dict": "object",
+    }
+
+    return type_mapping.get(type_name, "string")  # Default to string if unknown
+
+
 def client_tool(func: T) -> ClientTool:
     """
     Decorator to convert a function into a ClientTool.
@@ -188,13 +217,14 @@ def get_description(self) -> str:
                     f"No description found for client tool {__name__}. Please provide a RST-style docstring with description and :param tags for each parameter."
                 )
 
-        def get_params_definition(self) -> Dict[str, Parameter]:
+        def get_input_schema(self) -> JSONSchema:
             hints = get_type_hints(func)
             # Remove return annotation if present
             hints.pop("return", None)
 
             # Get parameter descriptions from docstring
-            params = {}
+            properties = {}
+            required = []
             sig = inspect.signature(func)
             doc = inspect.getdoc(func) or ""
 
@@ -212,15 +242,20 @@ def get_params_definition(self) -> Dict[str, Parameter]:
                 param = sig.parameters[name]
                 is_optional_type = get_origin(type_hint) is Union and type(None) in get_args(type_hint)
                 is_required = param.default == inspect.Parameter.empty and not is_optional_type
-                params[name] = Parameter(
-                    name=name,
-                    description=param_doc or f"Parameter {name}",
-                    parameter_type=type_hint.__name__,
-                    default=(param.default if param.default != inspect.Parameter.empty else None),
-                    required=is_required,
-                )
 
-            return params
+                properties[name] = {
+                    "type": _python_type_to_json_schema_type(type_hint),
+                    "description": param_doc,
+                }
+
+                if is_required:
+                    required.append(name)
+
+            return {
+                "type": "object",
+                "properties": properties,
+                "required": required,
+            }
 
         def run_impl(self, **kwargs) -> Any:
             if inspect.iscoroutinefunction(func):

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

@@ -37,7 +37,7 @@ def get_tool_defs(
                 {
                     "name": tool.identifier,
                     "description": tool.description,
-                    "parameters": tool.parameters,
+                    "input_schema": tool.input_schema,
                 }
                 for tool in client.tools.list(toolgroup_id=toolgroup_id)
             ]
@@ -48,7 +48,7 @@ def get_tool_defs(
             {
                 "name": tool.get_name(),
                 "description": tool.get_description(),
-                "parameters": tool.get_params_definition(),
+                "input_schema": tool.get_input_schema(),
             }
             for tool in client_tools
         ]