- Add support for enum tool params

- Add demo `arithmetic_calc.py`

Author: uogbuji

README.md

@@ -22,7 +22,7 @@ The following video, "Toolio in 10 minutes", is an easy way to learn about the p
 <img width="1268" alt="Toolio in 10 minutes still" src="https://github.com/user-attachments/assets/fc8dda94-326d-426e-a566-ac8ec60be31f">
 -->
 
-`toolio_server` is a FastAPI program that you can use to host MLX-format LLMs for structured output query, for example, if you are on  you can use the MLX format LLM model `mlx-community/Hermes-2-Theta-Llama-3-8B-4bit` as follows (from the cloned directory of this repository):
+`toolio_server` is a FastAPI program that you can use to host MLX-format LLMs for structured output query or function-calling. For example to host the MLX format LLM model `mlx-community/Hermes-2-Theta-Llama-3-8B-4bit` as follows (from the cloned directory of this repository):
 
 # Installation and Setup
 
@@ -33,7 +33,7 @@ toolio_server --model=mlx-community/Hermes-2-Theta-Llama-3-8B-4bit
 
 This will download the model (a little over 4GB) to your local HuggingFace disk cache, and running it will take up about that much of your unified RAM.
 
-For more on the MLX framework for ML workloads (including LLMs) on Apple Silicon, see the [MLX Notes](https://github.com/uogbuji/mlx-notes) article series. The "Day One" article provides all the context you need for using local LLMs through this project.
+For more on the MLX framework for ML workloads (including LLMs) on Apple Silicon, see the [MLX Notes](https://github.com/uogbuji/mlx-notes) article series. The "Day One" article provides all the context you need for using local LLMs with Toolio.
 
 # cURLing the Toolio server
 
@@ -91,7 +91,7 @@ toolio_request --apibase="http://localhost:8000" --prompt-file=/tmp/llmprompt.tx
 
 ## Tool calling
 
-You can also run tool usage (function-calling) prompts, a key technique in LLM agent frameworks. A schema will automatically be generated from the tool specs
+You can run tool usage (function-calling) prompts, a key technique in LLM agent frameworks. A schema will automatically be generated from the tool specs, which themselves are based on [JSON Schema](https://json-schema.org/), according to OpenAI conventions.
 
 ```sh
 echo 'What'\''s the weather like in Boulder today?' > /tmp/llmprompt.txt
@@ -389,6 +389,10 @@ async def query_sq_root(tmm):
 asyncio.run(query_sq_root(toolio_mm))
 ```
 
+# More examples
+
+See the `demo` directory.
+
 # Credits
 
 * otriscon's [llm-structured-output](https://github.com/otriscon/llm-structured-output/) is the foundation of this package

demo/arithmetic_calc.py

@@ -0,0 +1,56 @@
+'''
+demo/arithmetic_calc.py
+
+Based on an exaple from /llama-cpp-agent
+https://llama-cpp-agent.readthedocs.io/en/latest/parallel_function_calling/
+'''
+
+import asyncio
+from enum import Enum, auto
+
+from toolio.tool import tool, param
+from toolio.llm_helper import model_manager, extract_content
+
+
+class arithmetic_op(Enum):
+    ADD = 'add'
+    SUBTRACT = 'subtract'
+    MULTIPLY = 'multiply'
+    DIVIDE = 'divide'
+
+
+@tool('arithmetic_calc', params=[
+    param('num1', float, 'Number on the left hand side of the calculation', True),
+    param('num2', float, 'Number on the left hand side of the calculation', True),
+    param('op', arithmetic_op, 'Arithmetic operation to make on the two numbers', True),
+    ])
+async def arithmetic_calc(num1=None, num2=None, op=None):
+    'Very basic arithmetic calculator'
+    match op:
+        case arithmetic_op.ADD:
+            result = num1 + num2
+        case arithmetic_op.SUBTRACT:
+            result = num1 - num2
+        case arithmetic_op.MULTIPLY:
+            result = num1 * num2
+        case arithmetic_op.DIVIDE:
+            result = num1 / num2
+        case _:
+            raise ValueError('Unknown operator')  # Shouldn't happen
+    return result
+
+
+# Had a problem using Hermes-2-Theta-Llama-3-8B-4bit 😬
+# MLX_MODEL_PATH = 'mlx-community/Hermes-2-Theta-Llama-3-8B-4bit'
+MLX_MODEL_PATH = 'mlx-community/Mistral-Nemo-Instruct-2407-4bit'
+
+toolio_mm = model_manager(MLX_MODEL_PATH, tool_reg=[arithmetic_calc], trace=True)
+
+# PROMPT = 'Solve the following calculations: 42 * 42, 24 * 24, 5 * 5, 89 * 75, 42 * 46, 69 * 85, 422 * 420, 753 * 321, 72 * 55, 240 * 204, 789 * 654, 123 * 321, 432 * 89, 564 * 321?'  # noqa: E501
+PROMPT = 'Solve the following calculation: 4242 * 2424.2'
+async def async_main(tmm):
+    msgs = [ {'role': 'user', 'content': PROMPT} ]
+    async for chunk in extract_content(tmm.complete_with_tools(msgs)):
+        print(chunk, end='')
+
+asyncio.run(async_main(toolio_mm))

pylib/tool/schematics.py

@@ -9,6 +9,7 @@
 import functools
 import textwrap
 from dataclasses import dataclass
+from enum import Enum
 
 
 @dataclass
@@ -30,8 +31,15 @@ def tool_dec(func):
         required_list = []
         for p in params:
             # Translate type designation to JSON Schema, if need be
-            typ = TYPES_LOOKUP.get(p.typ, p.typ)
-            schema_params[p.name] = {'type': typ, 'description': p.desc}
+            typ = TYPES_LOOKUP.get(p.typ)
+            if typ:
+                schema_params[p.name] = {'type': typ, 'description': p.desc}
+            else:
+                typ = complex_type(p.typ)
+                if isinstance(typ, dict):
+                    schema_params[p.name] = {'description': p.desc, **typ}
+                else:
+                    raise RuntimeError('Unable to determine param type')
             params_lookup[p.name] = p
             if p.required:
                 required_list.append(p.name)
@@ -75,4 +83,16 @@ def tool_inner(*args, **kwargs):
     int: 'number',
     float: 'number',
     bool: 'boolean',
+    'string': 'string',
+    'number': 'number',
+    'boolean': 'boolean',
     }
+
+
+def complex_type(typ):
+    if issubclass(typ, Enum):
+        # Worth being aware of typ._member_names_ & typ._member_map_
+        return {'type': 'string', 'enum': [ e.value for e in typ ]}
+    else:
+        return None
+