01_funccall.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efe78920",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|default_exp funccall"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d773712-12fe-440e-891f-36f59666dfde",
   "metadata": {},
   "source": [
    "# funccall source"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5ad6b86",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "import inspect\n",
    "from collections import abc\n",
    "from fastcore.utils import *\n",
    "from fastcore.docments import docments\n",
    "from typing import get_origin, get_args, Dict, List, Optional, Tuple, Union\n",
    "from types import UnionType"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9f43047",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "empty = inspect.Parameter.empty"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a7cdbc6",
   "metadata": {},
   "source": [
    "## Function calling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ec35c95",
   "metadata": {},
   "source": [
    "Many LLMs do function calling (aka tool use) by taking advantage of JSON schema.\n",
    "\n",
    "We'll use [docments](https://fastcore.fast.ai/docments.html) to make getting JSON schema from Python functions as ergonomic as possible. Each parameter (and the return value) should have a type, and a docments comment with the description of what it is. Here's an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a017af1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def silly_sum(\n",
    "    a:int, # First thing to sum\n",
    "    b:int=1, # Second thing to sum\n",
    "    c:list[int]=None, # A pointless argument\n",
    ") -> int: # The sum of the inputs\n",
    "    \"Adds a + b.\"\n",
    "    return a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a3ff443",
   "metadata": {},
   "source": [
    "This is what `docments` makes of that:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3f2ebcf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "```json\n",
       "{ 'a': { 'anno': <class 'int'>,\n",
       "         'default': <class 'inspect._empty'>,\n",
       "         'docment': 'First thing to sum'},\n",
       "  'b': {'anno': <class 'int'>, 'default': 1, 'docment': 'Second thing to sum'},\n",
       "  'c': {'anno': list[int], 'default': None, 'docment': 'A pointless argument'},\n",
       "  'return': { 'anno': <class 'int'>,\n",
       "              'default': <class 'inspect._empty'>,\n",
       "              'docment': 'The sum of the inputs'}}\n",
       "```"
      ],
      "text/plain": [
       "{'a': {'docment': 'First thing to sum',\n",
       "  'anno': int,\n",
       "  'default': inspect._empty},\n",
       " 'b': {'docment': 'Second thing to sum', 'anno': int, 'default': 1},\n",
       " 'c': {'docment': 'A pointless argument', 'anno': list[int], 'default': None},\n",
       " 'return': {'docment': 'The sum of the inputs',\n",
       "  'anno': int,\n",
       "  'default': inspect._empty}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = docments(silly_sum, full=True)\n",
    "d"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "745e44ea",
   "metadata": {},
   "source": [
    "Note that this is an [AttrDict](https://fastcore.fast.ai/basics.html#attrdict) so we can treat it like an object, *or* a dict:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35cb279d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('First thing to sum', int)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d.a.docment, d['a']['anno']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7bf4025",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "def _types(t:type)->tuple[str,Optional[str]]:\n",
    "    \"Tuple of json schema type name and (if appropriate) array item name.\"\n",
    "    if t is empty: raise TypeError('Missing type')\n",
    "    tmap = {int:\"integer\", float:\"number\", str:\"string\", bool:\"boolean\", list:\"array\", dict:\"object\"}\n",
    "    tmap.update({k.__name__: v for k, v in tmap.items()})\n",
    "    if getattr(t, '__origin__', None) in (list,tuple):\n",
    "        args = getattr(t, '__args__', None)\n",
    "        item_type = \"object\" if not args else tmap.get(t.__args__[0].__name__, \"object\")\n",
    "        return \"array\", item_type\n",
    "    # if t is a string like 'int', directly use the string as the key\n",
    "    elif isinstance(t, str): return tmap.get(t, \"object\"), None\n",
    "    # if t is the type itself and a container\n",
    "    elif get_origin(t): return tmap.get(get_origin(t).__name__, \"object\"), None\n",
    "    # if t is the type itself like int, use the __name__ representation as the key\n",
    "    else: return tmap.get(t.__name__, \"object\"), None"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edf73046",
   "metadata": {},
   "source": [
    "This internal function is needed to convert Python types into JSON schema types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecb7bc52",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(('array', 'integer'), ('integer', None), ('integer', None))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_types(list[int]), _types(int), _types('int')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38b4650a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(('array', 'integer'), ('object', None), ('object', None), ('array', 'string'))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_types(List[int]), _types(Optional[str]), _types(str | None), _types(Tuple[str, int])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4d0ac1e",
   "metadata": {},
   "source": [
    "Note the current behavior:\n",
    "\n",
    "- ignores all but the first argument for tuples\n",
    "- union types map to object which is a stand-in for arbitrary types\n",
    "\n",
    "These and other approximations may require further refinement in the future."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0e3c940",
   "metadata": {},
   "source": [
    "Will also convert custom types to the `object` type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9969fd00",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(('array', 'object'), ('object', None))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Custom: a: int\n",
    "_types(list[Custom]), _types(Custom)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d5dc245",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "def _param(name, info):\n",
    "    \"json schema parameter given `name` and `info` from docments full dict.\"\n",
    "    paramt,itemt = _types(info.anno)\n",
    "    pschema = dict(type=paramt, description=info.docment or \"\")\n",
    "    if itemt: pschema[\"items\"] = {\"type\": itemt}\n",
    "    if info.default is not empty: pschema[\"default\"] = info.default\n",
    "    return pschema"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5337d6bd",
   "metadata": {},
   "source": [
    "This private function converts a key/value pair from the `docments` structure into the `dict` that will be needed for the schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2450ace6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a // {'docment': 'First thing to sum', 'anno': <class 'int'>, 'default': <class 'inspect._empty'>}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'type': 'integer', 'description': 'First thing to sum'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n,o = first(d.items())\n",
    "print(n,'//', o)\n",
    "_param(n, o)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba6bcac4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "custom_types = {Path}\n",
    "\n",
    "def _handle_type(t, defs):\n",
    "    \"Handle a single type, creating nested schemas if necessary\"\n",
    "    if t is NoneType: return {'type': 'null'}\n",
    "    if t in custom_types: return {'type':'string', 'format':t.__name__}\n",
    "    if isinstance(t, type) and not issubclass(t, (int, float, str, bool)) or inspect.isfunction(t):\n",
    "        defs[t.__name__] = _get_nested_schema(t)\n",
    "        return {'$ref': f'#/$defs/{t.__name__}'}\n",
    "    return {'type': _types(t)[0]}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16dbf080",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "({'type': 'integer'}, {'type': 'string', 'format': 'Path'})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_handle_type(int, None), _handle_type(Path, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fd6cd29",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _is_container(t):\n",
    "    \"Check if type is a container (list, dict, tuple, set, Union)\"\n",
    "    origin = get_origin(t)\n",
    "    return origin in (list, dict, tuple, set, Union) if origin else False\n",
    "\n",
    "def _is_parameterized(t):\n",
    "    \"Check if type has arguments (e.g. list[int] vs list, dict[str, int] vs dict)\"\n",
    "    return _is_container(t) and (get_args(t) != ())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "783747af",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert _is_parameterized(list[int]) == True\n",
    "assert _is_parameterized(int) == False\n",
    "assert _is_container(list[int]) == True\n",
    "assert _is_container(dict[str, int]) == True\n",
    "assert _is_container(int) == False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d42c88dd",
   "metadata": {},
   "source": [
    "For union and optional types, `Union` covers older `Union[str]` syntax while `UnionType` covers 3.10+ `str | None` syntax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7815799b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(str | None, types.UnionType, (str, NoneType))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _example_new_unioin(opt_tup: str | None):\n",
    "    pass\n",
    "\n",
    "d = docments(_example_new_unioin, full=True)\n",
    "anno1 = first(d.items())[1].anno\n",
    "(anno1, get_origin(anno1), get_args(anno1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d745c902",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(typing.Optional[str], typing.Union, (str, NoneType))"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _example_old_union(opt_tup: Union[str, type(None)] =None):\n",
    "    pass\n",
    "\n",
    "d = docments(_example_old_union, full=True)\n",
    "anno2 = first(d.items())[1].anno\n",
    "(anno2, get_origin(anno2), get_args(anno2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c5701c7",
   "metadata": {},
   "source": [
    "Support for both union types is part of the broader container handling:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1153f02",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _handle_container(origin, args, defs):\n",
    "    \"Handle container types like dict, list, tuple, set, and Union\"\n",
    "    if origin is Union or origin is UnionType:\n",
    "        return {\"anyOf\": [_handle_type(arg, defs) for arg in args]}\n",
    "    if origin is dict:\n",
    "        value_type = args[1].__args__[0] if hasattr(args[1], '__args__') else args[1]\n",
    "        return {\n",
    "            'type': 'object',\n",
    "            'additionalProperties': (\n",
    "                {'type': 'array', 'items': _handle_type(value_type, defs)}\n",
    "                if hasattr(args[1], '__origin__') else _handle_type(args[1], defs)\n",
    "            )\n",
    "        }\n",
    "    elif origin in (list, tuple, set):\n",
    "        schema = {'type': 'array', 'items': _handle_type(args[0], defs)}\n",
    "        if origin is set:\n",
    "            schema['uniqueItems'] = True\n",
    "        return schema\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ee1c529",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _process_property(name, obj, props, req, defs):\n",
    "    \"Process a single property of the schema\"\n",
    "    p = _param(name, obj)\n",
    "    props[name] = p\n",
    "    if obj.default is empty: req[name] = True\n",
    "\n",
    "    if _is_container(obj.anno) and _is_parameterized(obj.anno):\n",
    "            p.update(_handle_container(get_origin(obj.anno), get_args(obj.anno), defs))        \n",
    "    else:\n",
    "        # Non-container type or container without arguments\n",
    "        p.update(_handle_type(obj.anno, defs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38b0f97e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def _get_nested_schema(obj):\n",
    "    \"Generate nested JSON schema for a class or function\"\n",
    "    d = docments(obj, full=True)\n",
    "    props, req, defs = {}, {}, {}\n",
    "\n",
    "    for n, o in d.items():\n",
    "        if n != 'return' and n != 'self':\n",
    "            _process_property(n, o, props, req, defs)\n",
    "\n",
    "    schema = dict(type='object', properties=props, title=obj.__name__ if isinstance(obj, type) else None)\n",
    "    if req: schema['required'] = list(req)\n",
    "    if defs: schema['$defs'] = defs\n",
    "    return schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bb9df6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test primitive types\n",
    "defs = {}\n",
    "assert _handle_type(int, defs) == {'type': 'integer'}\n",
    "assert _handle_type(str, defs) == {'type': 'string'}\n",
    "assert _handle_type(bool, defs) == {'type': 'boolean'}\n",
    "assert _handle_type(float, defs) == {'type': 'number'}\n",
    "\n",
    "# Test custom class\n",
    "class TestClass:\n",
    "    def __init__(self, x: int, y: int): store_attr()\n",
    "\n",
    "result = _handle_type(TestClass, defs)\n",
    "assert result == {'$ref': '#/$defs/TestClass'}\n",
    "assert 'TestClass' in defs\n",
    "assert defs['TestClass']['type'] == 'object'\n",
    "assert 'properties' in defs['TestClass']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1d09435",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test primitive types in containers\n",
    "assert _handle_container(list, (int,), defs) == {'type': 'array', 'items': {'type': 'integer'}}\n",
    "assert _handle_container(tuple, (str,), defs) == {'type': 'array', 'items': {'type': 'string'}}\n",
    "assert _handle_container(set, (str,), defs) == {'type': 'array', 'items': {'type': 'string'}, 'uniqueItems': True}\n",
    "assert _handle_container(dict, (str,bool), defs) == {'type': 'object', 'additionalProperties': {'type': 'boolean'}}\n",
    "\n",
    "result = _handle_container(list, (TestClass,), defs)\n",
    "assert result == {'type': 'array', 'items': {'$ref': '#/$defs/TestClass'}}\n",
    "assert 'TestClass' in defs\n",
    "\n",
    "# Test complex nested structure\n",
    "ComplexType = dict[str, list[TestClass]]\n",
    "result = _handle_container(dict, (str, list[TestClass]), defs)\n",
    "assert result == {\n",
    "    'type': 'object',\n",
    "    'additionalProperties': {\n",
    "        'type': 'array',\n",
    "        'items': {'$ref': '#/$defs/TestClass'}\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5fd37d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test processing of a required integer property\n",
    "props, req = {}, {}\n",
    "class TestClass:\n",
    "    \"Test class\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        x: int, # First thing\n",
    "        y: list[float], # Second thing\n",
    "        z: str = \"default\", # Third thing\n",
    "    ): store_attr()\n",
    "\n",
    "d = docments(TestClass, full=True)\n",
    "_process_property('x', d.x, props, req, defs)\n",
    "assert 'x' in props\n",
    "assert props['x']['type'] == 'integer'\n",
    "assert 'x' in req\n",
    "\n",
    "# Test processing of a required list property\n",
    "_process_property('y', d.y, props, req, defs)\n",
    "assert 'y' in props\n",
    "assert props['y']['type'] == 'array'\n",
    "assert props['y']['items']['type'] == 'number'\n",
    "assert 'y' in req\n",
    "\n",
    "# Test processing of an optional string property with default\n",
    "_process_property('z', d.z, props, req, defs)\n",
    "assert 'z' in props\n",
    "assert props['z']['type'] == 'string'\n",
    "assert props['z']['default'] == \"default\"\n",
    "assert 'z' not in req"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23f54386",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "def get_schema(f:Union[callable,dict], pname='input_schema')->dict:\n",
    "    \"Generate JSON schema for a class, function, or method\"\n",
    "    if isinstance(f, dict): return f\n",
    "    schema = _get_nested_schema(f)\n",
    "    desc = f.__doc__\n",
    "    assert desc, \"Docstring missing!\"\n",
    "    d = docments(f, full=True)\n",
    "    ret = d.pop('return')\n",
    "    if ret.anno is not empty: desc += f'\\n\\nReturns:\\n- type: {_types(ret.anno)[0]}'\n",
    "    return {\"name\": f.__name__, \"description\": desc, pname: schema}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a59df671",
   "metadata": {},
   "source": [
    "Putting this all together, we can now test getting a schema from `silly_sum`. The tool use spec doesn't support return annotations directly, so we put that in the description instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7311af9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adds a + b.\n",
      "\n",
      "Returns:\n",
      "- type: integer\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'name': 'silly_sum',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'a': {'type': 'integer', 'description': 'First thing to sum'},\n",
       "   'b': {'type': 'integer',\n",
       "    'description': 'Second thing to sum',\n",
       "    'default': 1},\n",
       "   'c': {'type': 'array',\n",
       "    'description': 'A pointless argument',\n",
       "    'items': {'type': 'integer'},\n",
       "    'default': None}},\n",
       "  'title': None,\n",
       "  'required': ['a']}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s = get_schema(silly_sum)\n",
    "desc = s.pop('description')\n",
    "print(desc)\n",
    "s"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d478ba6b",
   "metadata": {},
   "source": [
    "This also works with string annotations, e.g:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80203962",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'silly_test',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'a': {'type': 'integer', 'description': 'quoted type hint'}},\n",
       "  'title': None,\n",
       "  'required': ['a']}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def silly_test(\n",
    "    a: 'int',  # quoted type hint\n",
    "):\n",
    "    \"Mandatory docstring\"\n",
    "    return a\n",
    "\n",
    "get_schema(silly_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3f36f8a",
   "metadata": {},
   "source": [
    "This also works with instance methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05d33447",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'sums',\n",
       " 'description': 'Adds a + b.\\n\\nReturns:\\n- type: integer',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'a': {'type': 'integer', 'description': 'First thing to sum'},\n",
       "   'b': {'type': 'integer',\n",
       "    'description': 'Second thing to sum',\n",
       "    'default': 1}},\n",
       "  'title': None,\n",
       "  'required': ['a']}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Dummy:\n",
    "    def sums(\n",
    "        self,\n",
    "        a:int,  # First thing to sum\n",
    "        b:int=1 # Second thing to sum\n",
    "    ) -> int: # The sum of the inputs\n",
    "        \"Adds a + b.\"\n",
    "        print(f\"Finding the sum of {a} and {b}\")\n",
    "        return a + b\n",
    "\n",
    "get_schema(Dummy.sums)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae3fdfa4",
   "metadata": {},
   "source": [
    "`get_schema` also handles more complicated structures such as nested classes. This is useful for things like structured outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce3be915",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'Conversation',\n",
       " 'description': 'A conversation between two speakers',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'turns': {'type': 'array',\n",
       "    'description': 'Turns of the conversation',\n",
       "    'items': {'$ref': '#/$defs/Turn'}}},\n",
       "  'title': 'Conversation',\n",
       "  'required': ['turns'],\n",
       "  '$defs': {'Turn': {'type': 'object',\n",
       "    'properties': {'speaker_a': {'type': 'string',\n",
       "      'description': \"First speaker's message\"},\n",
       "     'speaker_b': {'type': 'string',\n",
       "      'description': \"Second speaker's message\"}},\n",
       "    'title': 'Turn',\n",
       "    'required': ['speaker_a', 'speaker_b']}}}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Turn:\n",
    "    \"Turn between two speakers\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        speaker_a:str, # First speaker's message\n",
    "        speaker_b:str,  # Second speaker's message\n",
    "    ): store_attr()\n",
    "\n",
    "class Conversation:\n",
    "    \"A conversation between two speakers\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        turns:list[Turn], # Turns of the conversation\n",
    "    ): store_attr()\n",
    "\n",
    "get_schema(Conversation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "386e514d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'DictConversation',\n",
       " 'description': 'A conversation between two speakers',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'turns': {'type': 'object',\n",
       "    'description': 'dictionary of topics and the Turns of the conversation',\n",
       "    'additionalProperties': {'type': 'array',\n",
       "     'items': {'$ref': '#/$defs/Turn'}}}},\n",
       "  'title': 'DictConversation',\n",
       "  'required': ['turns'],\n",
       "  '$defs': {'Turn': {'type': 'object',\n",
       "    'properties': {'speaker_a': {'type': 'string',\n",
       "      'description': \"First speaker's message\"},\n",
       "     'speaker_b': {'type': 'string',\n",
       "      'description': \"Second speaker's message\"}},\n",
       "    'title': 'Turn',\n",
       "    'required': ['speaker_a', 'speaker_b']}}}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class DictConversation:\n",
    "    \"A conversation between two speakers\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        turns:dict[str,list[Turn]], # dictionary of topics and the Turns of the conversation\n",
    "    ): store_attr()\n",
    "\n",
    "get_schema(DictConversation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c08ac6b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'SetConversation',\n",
       " 'description': 'A conversation between two speakers',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'turns': {'type': 'array',\n",
       "    'description': 'the unique Turns of the conversation',\n",
       "    'items': {'$ref': '#/$defs/Turn'},\n",
       "    'uniqueItems': True}},\n",
       "  'title': 'SetConversation',\n",
       "  'required': ['turns'],\n",
       "  '$defs': {'Turn': {'type': 'object',\n",
       "    'properties': {'speaker_a': {'type': 'string',\n",
       "      'description': \"First speaker's message\"},\n",
       "     'speaker_b': {'type': 'string',\n",
       "      'description': \"Second speaker's message\"}},\n",
       "    'title': 'Turn',\n",
       "    'required': ['speaker_a', 'speaker_b']}}}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class SetConversation:\n",
    "    \"A conversation between two speakers\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        turns:set[Turn], # the unique Turns of the conversation\n",
    "    ): store_attr()\n",
    "\n",
    "get_schema(SetConversation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cf3f35c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "def PathArg(\n",
    "    path: str  # A filesystem path\n",
    "): return Path(path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "169212a6",
   "metadata": {},
   "source": [
    "Paths are a special case, since they only take `*args` and `**kwargs` as params, but normally we'd use them in a schema by just passing a str. So we create a custom param type for that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9135dfa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'path_test',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'a': {'type': 'object',\n",
       "    'description': 'a type hint',\n",
       "    '$ref': '#/$defs/PathArg'},\n",
       "   'b': {'type': 'object',\n",
       "    'description': 'b type hint',\n",
       "    '$ref': '#/$defs/PathArg'}},\n",
       "  'title': None,\n",
       "  'required': ['a', 'b'],\n",
       "  '$defs': {'PathArg': {'type': 'object',\n",
       "    'properties': {'path': {'type': 'string',\n",
       "      'description': 'A filesystem path'}},\n",
       "    'title': None,\n",
       "    'required': ['path']}}}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def path_test(\n",
    "    a: PathArg,  # a type hint\n",
    "    b: PathArg   # b type hint\n",
    "):\n",
    "    \"Mandatory docstring\"\n",
    "    return a/b\n",
    "\n",
    "get_schema(path_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6d1d0c8",
   "metadata": {},
   "source": [
    "Alternatively, use `Path` as usual, and handle the `format` key in the json to use that as a callable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdb69462",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': 'path_test2',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'a': {'type': 'string',\n",
       "    'description': 'a type hint',\n",
       "    'format': 'Path'},\n",
       "   'b': {'type': 'string', 'description': 'b type hint', 'format': 'Path'}},\n",
       "  'title': None,\n",
       "  'required': ['a', 'b']}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def path_test2(\n",
    "    a: Path,  # a type hint\n",
    "    b: Path   # b type hint\n",
    "):\n",
    "    \"Mandatory docstring\"\n",
    "    return a/b\n",
    "\n",
    "get_schema(path_test2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "369320d4",
   "metadata": {},
   "source": [
    "### Additional `get_schema()` Test Cases"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8052380",
   "metadata": {},
   "source": [
    "Union types are approximately mapped to JSON schema 'anyOf' with two or more value types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fc1d6f9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': '_union_test',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'opt_tup': {'type': 'object',\n",
       "    'description': '',\n",
       "    'default': None,\n",
       "    'anyOf': [{'type': 'array'}, {'type': 'string'}, {'type': 'integer'}]}},\n",
       "  'title': None}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _union_test(opt_tup: Union[Tuple[int, int], str, int]=None):\n",
    "    \"Mandatory docstring\"\n",
    "    return \"\"\n",
    "get_schema(_union_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7641aca8",
   "metadata": {},
   "source": [
    "The new (Python 3.10+) union syntax can also be used, producing an equivalent schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1a11b3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': '_new_union_test',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'opt_tup': {'type': 'object',\n",
       "    'description': '',\n",
       "    'default': None,\n",
       "    'anyOf': [{'type': 'array'}, {'type': 'string'}, {'type': 'integer'}]}},\n",
       "  'title': None}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _new_union_test(opt_tup: Tuple[int, int] | str | int =None):\n",
    "    \"Mandatory docstring\"\n",
    "    pass\n",
    "get_schema(_new_union_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d24cc0a",
   "metadata": {},
   "source": [
    "Optional is a special case of union types, limited to two types, one of which is None (mapped to null in JSON schema):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac8f3d19",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': '_optional_test',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'opt_tup': {'type': 'object',\n",
       "    'description': '',\n",
       "    'default': None,\n",
       "    'anyOf': [{'type': 'array'}, {'type': 'null'}]}},\n",
       "  'title': None}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _optional_test(opt_tup: Optional[Tuple[int, int]]=None):\n",
    "    \"Mandatory docstring\"\n",
    "    pass\n",
    "get_schema(_optional_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c969721b",
   "metadata": {},
   "source": [
    "Containers can also be used, both in their parameterized form (`List[int]`) or as their unparameterized raw type (`List`). In the latter case, the item type is mapped to `object` in JSON schema."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2959197",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': '_list_test',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'l': {'type': 'array',\n",
       "    'description': '',\n",
       "    'items': {'type': 'integer'}}},\n",
       "  'title': None,\n",
       "  'required': ['l']}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _list_test(l: List[int]):\n",
    "    \"Mandatory docstring\"\n",
    "    pass\n",
    "get_schema(_list_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8fbfea7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': '_raw_list_test',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'l': {'type': 'array',\n",
       "    'description': '',\n",
       "    'items': {'type': 'object'}}},\n",
       "  'title': None,\n",
       "  'required': ['l']}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _raw_list_test(l: List):\n",
    "    \"Mandatory docstring\"\n",
    "    pass\n",
    "get_schema(_raw_list_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5704c197",
   "metadata": {},
   "source": [
    "The same applies to dictionary, which can similarly be parameterized with key/value types or specified as a raw type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2e8c567",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': '_dict_test',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'d': {'type': 'object',\n",
       "    'description': '',\n",
       "    'additionalProperties': {'type': 'integer'}}},\n",
       "  'title': None,\n",
       "  'required': ['d']}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _dict_test(d: Dict[str, int]):\n",
    "    \"Mandatory docstring\"\n",
    "    pass\n",
    "get_schema(_dict_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3138ac4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'name': '_raw_dict_test',\n",
       " 'description': 'Mandatory docstring',\n",
       " 'input_schema': {'type': 'object',\n",
       "  'properties': {'d': {'type': 'object', 'description': ''}},\n",
       "  'title': None,\n",
       "  'required': ['d']}}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def _raw_dict_test(d: Dict):\n",
    "    \"Mandatory docstring\"\n",
    "get_schema(_raw_dict_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9529d39a",
   "metadata": {},
   "source": [
    "### Python tool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a69cad9",
   "metadata": {},
   "source": [
    "In language model clients it's often useful to have a 'code interpreter' -- this is something that runs code, and generally outputs the result of the last expression (i.e like IPython or Jupyter). \n",
    "\n",
    "In this section we'll create the `python` function, which executes a string as Python code, with an optional timeout. If the last line is an expression, we'll return that -- just like in IPython or Jupyter, but without needing them installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873000d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "import ast, time, signal, traceback\n",
    "from fastcore.utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4703296a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "def _copy_loc(new, orig):\n",
    "    \"Copy location information from original node to new node and all children.\"\n",
    "    new = ast.copy_location(new, orig)\n",
    "    for field, o in ast.iter_fields(new):\n",
    "        if isinstance(o, ast.AST): setattr(new, field, _copy_loc(o, orig))\n",
    "        elif isinstance(o, list): setattr(new, field, [_copy_loc(value, orig) for value in o])\n",
    "    return new"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c0d4922",
   "metadata": {},
   "source": [
    "This is an internal function that's needed for `_run` to ensure that location information is available in the abstract syntax tree (AST), since otherwise python complains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1574585f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "def _run(code:str, glb:dict=None, loc:dict=None):\n",
    "    \"Run `code`, returning final expression (similar to IPython)\"\n",
    "    tree = ast.parse(code)\n",
    "    last_node = tree.body[-1] if tree.body else None\n",
    "    \n",
    "    # If the last node is an expression, modify the AST to capture the result\n",
    "    if isinstance(last_node, ast.Expr):\n",
    "        tgt = [ast.Name(id='_result', ctx=ast.Store())]\n",
    "        assign_node = ast.Assign(targets=tgt, value=last_node.value)\n",
    "        tree.body[-1] = _copy_loc(assign_node, last_node)\n",
    "\n",
    "    compiled_code = compile(tree, filename='<ast>', mode='exec')\n",
    "    glb = glb or {}\n",
    "    stdout_buffer = io.StringIO()\n",
    "    saved_stdout = sys.stdout\n",
    "    sys.stdout = stdout_buffer\n",
    "    try: exec(compiled_code, glb, loc)\n",
    "    finally: sys.stdout = saved_stdout\n",
    "    _result = glb.get('_result', None)\n",
    "    if _result is not None: return _result\n",
    "    return stdout_buffer.getvalue().strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92ca7f47",
   "metadata": {},
   "source": [
    "This is the internal function used to actually run the code -- we pull off the last AST to see if it's an expression (i.e something that returns a value), and if so, we store it to a special `_result` variable so we can return it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15b72cb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "479001600"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_run('import math;math.factorial(12)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "632a7ac1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2'"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "_run('print(1+1)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34f2e5c2",
   "metadata": {},
   "source": [
    "We now have the machinery needed to create our `python` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81857615",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "def python(code:str, # Code to execute\n",
    "           glb:Optional[dict]=None, # Globals namespace\n",
    "           loc:Optional[dict]=None, # Locals namespace\n",
    "           timeout:int=3600 # Maximum run time in seconds before a `TimeoutError` is raised\n",
    "          ): # Result of last node, if it's an expression, or `None` otherwise\n",
    "    \"\"\"Executes python `code` with `timeout` and returning final expression (similar to IPython).\n",
    "    Raised exceptions are returned as a string, with a stack trace.\"\"\"\n",
    "    def handler(*args): raise TimeoutError()\n",
    "    if glb is None: glb = inspect.currentframe().f_back.f_globals\n",
    "    if loc is None: loc=glb\n",
    "    signal.signal(signal.SIGALRM, handler)\n",
    "    signal.alarm(timeout)\n",
    "    try: return _run(code, glb, loc)\n",
    "    except Exception as e: return traceback.format_exc()\n",
    "    finally: signal.alarm(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b9324f",
   "metadata": {},
   "source": [
    "There's no builtin security here -- you should generally use this in a sandbox, or alternatively prompt before running code. It can handle multiline function definitions, and pretty much any other normal Python syntax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69d74f4d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "120"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "python(\"\"\"def factorial(n):\n",
    "    if n == 0 or n == 1: return 1\n",
    "    else: return n * factorial(n-1)\n",
    "factorial(5)\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c629442",
   "metadata": {},
   "source": [
    "If the code takes longer than `timeout` then it returns an error string."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcb472b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Traceback (most recent call last):\n",
      "  File \"/var/folders/51/b2_szf2945n072c0vj2cyty40000gn/T/ipykernel_46879/2963369439.py\", line 14, in python\n",
      "    try: return _run(code, glb, loc)\n",
      "                ^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"/var/folders/51/b2_szf2945n072c0vj2cyty40000gn/T/ipykernel_46879/1858893181.py\", line 18, in _run\n",
      "    try: exec(compiled_code, glb, loc)\n",
      "         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n",
      "  File \"<ast>\", line 1, in <module>\n",
      "  File \"/var/folders/51/b2_szf2945n072c0vj2cyty40000gn/T/ipykernel_46879/2963369439.py\", line 9, in handler\n",
      "    def handler(*args): raise TimeoutError()\n",
      "                        ^^^^^^^^^^^^^^^^^^^^\n",
      "TimeoutError\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(python('import time; time.sleep(10)', timeout=1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d45684c1",
   "metadata": {},
   "source": [
    "By default the caller's global namespace is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72dfe290",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "python(\"a=1\")\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf48557c",
   "metadata": {},
   "source": [
    "Pass a different `glb` if needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55fb5613",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 3)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "glb = {}\n",
    "python(\"a=3\", glb)\n",
    "a, glb['a']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "244c502e",
   "metadata": {},
   "source": [
    "### Tool Calling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "186408f8",
   "metadata": {},
   "source": [
    "Many LLM API providers offer tool calling where an LLM can choose to call a given tool. This is also helpful for structured outputs since the response from the LLM is contrained to the required arguments of the tool.\n",
    "\n",
    "This section will be dedicated to helper functions for calling tools. We don't want to allow LLMs to call just any possible function (that would be a security disaster!) so we create a namespace -- that is, a dictionary of allowable function names to call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "782c4415",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def mk_ns(*funcs_or_objs):\n",
    "    merged = {}\n",
    "    for o in funcs_or_objs:\n",
    "        if isinstance(o, type): merged |= {n:getattr(o,n) for n,m in o.__dict__.items() if isinstance(m, (staticmethod, classmethod))}\n",
    "        if isinstance(o, object): merged |= {n:getattr(o,n) for n, m in inspect.getmembers(o, inspect.ismethod)} | {n:m for n,m in o.__class__.__dict__.items() if isinstance(m, staticmethod)}\n",
    "        if callable(o) and hasattr(o, '__name__'): merged |= {o.__name__: o}\n",
    "    return merged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5947aac4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'sums': <function __main__.sums(a, b)>}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def sums(a, b): return a + b\n",
    "ns = mk_ns(sums); ns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86ce0458",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ns['sums'](1, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d22f82",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dummy:\n",
    "    def __init__(self,a): self.a = a\n",
    "    def __call__(self): return self.a\n",
    "    def sums(self, a, b): return a + b\n",
    "    @staticmethod\n",
    "    def subs(a, b): return a - b\n",
    "    @classmethod\n",
    "    def mults(cls, a, b): return a * b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca50b957",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'subs': <function __main__.Dummy.subs(a, b)>,\n",
       " 'mults': <bound method Dummy.mults of <class '__main__.Dummy'>>,\n",
       " 'Dummy': __main__.Dummy}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ns = mk_ns(Dummy); ns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59ef734f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-1, 6)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ns['subs'](1, 2), ns['mults'](3, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15871e6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'__call__': <bound method Dummy.__call__ of <__main__.Dummy object>>,\n",
       " '__init__': <bound method Dummy.__init__ of <__main__.Dummy object>>,\n",
       " 'mults': <bound method Dummy.mults of <class '__main__.Dummy'>>,\n",
       " 'sums': <bound method Dummy.sums of <__main__.Dummy object>>,\n",
       " 'subs': <staticmethod(<function Dummy.subs>)>}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = Dummy(10)\n",
    "ns = mk_ns(d); ns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13cb7685",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-1, 6, 5, 10)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ns['subs'](1, 2), ns['mults'](3, 2), ns['sums'](3, 2), ns['__call__']()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2dfe13ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(None, -99)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ns['__init__'](-99), ns['__call__']()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85b4734f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "def call_func(fc_name, fc_inputs, ns):\n",
    "    \"Call the function `fc_name` with the given `fc_inputs` using namespace `ns`.\"\n",
    "    if not isinstance(ns, abc.Mapping): ns = mk_ns(*ns)\n",
    "    func = ns[fc_name]\n",
    "    return func(**fc_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce9cce60",
   "metadata": {},
   "source": [
    "Now when we an LLM responses with the tool to use and its inputs, we can simply use the same namespace it was given to look up the tool and call it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2ade8a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await call_func('sums', {'a': 1, 'b': 2}, ns=[sums])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aace64a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-1"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await call_func('subs', {'a': 1, 'b': 2}, ns=mk_ns(d))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "591574b8-6b53-4908-8159-b87be42133f7",
   "metadata": {},
   "source": [
    "### Async function calling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96a3a7d3-31ef-4cc6-b47c-35eaa8bbff8b",
   "metadata": {},
   "source": [
    "Since tools defined by MCP servers are async function, it is probably a good idea to have an async version of `call_func`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e273507b-6e4b-40bb-ae23-6397e89a4d51",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'asums': <function __main__.asums(a, b)>}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "async def asums(a, b): return a + b\n",
    "ns = mk_ns(asums); ns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ac04e80-7bb9-4b52-8285-454684605d47",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exports\n",
    "async def call_func_async(fc_name, fc_inputs, ns):\n",
    "    \"Awaits the function `fc_name` with the given `fc_inputs` using namespace `ns`.\"\n",
    "    if not isinstance(ns, abc.Mapping): ns = mk_ns(*ns)\n",
    "    func = ns[fc_name]\n",
    "    res = func(**fc_inputs)\n",
    "    if inspect.iscoroutine(res): res = await res\n",
    "    return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b83998ac-68e2-4dbe-b594-65fb4fdf59b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "await call_func_async('asums', {'a': 1, 'b': 2}, ns=[asums])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94ec4289",
   "metadata": {},
   "source": [
    "## Export -"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e9ee5c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "#|eval: false\n",
    "from nbdev.doclinks import nbdev_export\n",
    "nbdev_export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "207f9715",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}