notebooks/1.4-sr-ollama_function_calling.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/suriya/projects/dspy\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/suriya/projects/dspy/env-dspy/lib/python3.10/site-packages/IPython/core/magics/osm.py:417: UserWarning: This is now an optional IPython functionality, setting dhist requires you to install the `pickleshare` library.\n",
      "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n"
     ]
    }
   ],
   "source": [
    "cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tool Support Example\n",
    "\n",
    "[tool-support](https://ollama.com/blog/tool-support)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'function': {'name': 'get_current_weather', 'arguments': {'city': 'Toronto'}}}]\n"
     ]
    }
   ],
   "source": [
    "import ollama\n",
    "\n",
    "response = ollama.chat(\n",
    "    model='llama3.1',\n",
    "    messages=[{'role': 'user', 'content': \n",
    "        'What is the weather in Toronto?'}],\n",
    "\n",
    "\t\t# provide a weather checking tool to the model\n",
    "    tools=[{\n",
    "      'type': 'function',\n",
    "      'function': {\n",
    "        'name': 'get_current_weather',\n",
    "        'description': 'Get the current weather for a city',\n",
    "        'parameters': {\n",
    "          'type': 'object',\n",
    "          'properties': {\n",
    "            'city': {\n",
    "              'type': 'string',\n",
    "              'description': 'The name of the city',\n",
    "            },\n",
    "          },\n",
    "          'required': ['city'],\n",
    "        },\n",
    "      },\n",
    "    },\n",
    "  ],\n",
    ")\n",
    "\n",
    "print(response['message']['tool_calls'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A real example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[*********************100%***********************]  1 of 1 completed\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "207.22999572753906"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import yfinance as yf\n",
    "TICKER = \"AAPL\"\n",
    "\n",
    "def getHistoricalPrice(start: str, end: str):\n",
    "    df = yf.download(TICKER, start, end)\n",
    "    return df[\"Close\"]\n",
    "\n",
    "\n",
    "def computeStatistic(start: str, end: str, statistic: str=\"mean\"):\n",
    "    data = getHistoricalPrice(start, end)\n",
    "    return {\n",
    "        \"mean\": data.mean(),\n",
    "        \"max\": data.max(),\n",
    "        \"min\": data.min()\n",
    "    }[statistic].item()\n",
    "\n",
    "start, end = \"2024-08-01\", \"2024-08-31\"\n",
    "computeStatistic(start, end, statistic=\"min\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Describe the function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [{\n",
    "        'type': 'function',\n",
    "        'function': {\n",
    "            'name': 'getHistoricalPrice',\n",
    "            'description': '''\n",
    "            Get historical stock price given start and end dates.\n",
    "            Example: Give me last month\\'s stock price data.\n",
    "            '''\n",
    "            ,\n",
    "            'parameters': {\n",
    "            'type': 'object',\n",
    "            'properties': {\n",
    "                'start': {\n",
    "                'type': 'string',\n",
    "                'description': 'Start Date in YYYY-MM-DD format',\n",
    "                },\n",
    "                'end': {\n",
    "                'type': 'string',\n",
    "                'description': 'End Date in YYYY-MM-DD format',\n",
    "                },\n",
    "            },\n",
    "            'required': ['start', 'end'],\n",
    "            },\n",
    "        },\n",
    "        },\n",
    "        {\n",
    "        'type': 'function',\n",
    "        'function': {\n",
    "            'name': 'computeStatistic',\n",
    "            'description': '''\n",
    "            Given start and end dates compute a summary statistic represented by statistic argument.\n",
    "            Example: \"What is the maximum stock price this year?\".\n",
    "            '''\n",
    "            ,\n",
    "            'parameters': {\n",
    "            'type': 'object',\n",
    "            'properties': {\n",
    "                'start': {\n",
    "                'type': 'string',\n",
    "                'description': 'Start Date in YYYY-MM-DD format',\n",
    "                },\n",
    "                'end': {\n",
    "                'type': 'string',\n",
    "                'description': 'End Date in YYYY-MM-DD format',\n",
    "                },\n",
    "                'statistic': {\n",
    "                'type': 'string',\n",
    "                'description': 'Summary statistic to compute. Options include \"min\", \"max\", and, \"mean\".',\n",
    "                },\n",
    "            },\n",
    "            'required': ['start', 'end', 'statistic'],\n",
    "            },\n",
    "        },\n",
    "        },\n",
    "    ]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "DATE_FORMAT = \"%Y-%m-%d\"\n",
    "\n",
    "\n",
    "def today():\n",
    "    return datetime.now().strftime(DATE_FORMAT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 .  Give me this year's stock price data.\n",
      "[{'function': {'name': 'getHistoricalPrice', 'arguments': {'end': '2024-09-02', 'start': '2024-01-01'}}}]\n",
      "---\n",
      "2 .  Lowest price in 2023\n",
      "[{'function': {'name': 'computeStatistic', 'arguments': {'end': '2023-12-31', 'start': '2023-01-01', 'statistic': 'min'}}}]\n",
      "---\n"
     ]
    }
   ],
   "source": [
    "import ollama\n",
    "\n",
    "\n",
    "def generate_response(query):\n",
    "    response = ollama.chat(\n",
    "        model='llama3.1',\n",
    "        messages=[{\"role\": \"system\", \"content\": f\"today's date is {today()}\"}, {'role': 'user', 'content': query}],\n",
    "        tools=tools\n",
    "    )\n",
    "    content = response[\"message\"][\"content\"]\n",
    "    if content:\n",
    "        return content\n",
    "    return response[\"message\"][\"tool_calls\"]\n",
    "\n",
    "queries = [\n",
    "    \"Give me this year's stock price data.\",\n",
    "    \"Lowest price in 2023\",\n",
    "]\n",
    "for i,q in enumerate(queries):\n",
    "    print(i+1, \". \", q)\n",
    "    print(generate_response(q))\n",
    "    print(\"---\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That worked out great for us. \n",
    "Now, let's put it into action."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function calling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "\n",
    "def call_function(fn):\n",
    "    args = fn[\"function\"][\"arguments\"]\n",
    "    name = fn[\"function\"][\"name\"]\n",
    "\n",
    "    return {\n",
    "        \"getHistoricalPrice\": getHistoricalPrice,\n",
    "        \"computeStatistic\": computeStatistic\n",
    "    }[name](**args)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rewrite `generate_response` to accomodate function calling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_response(query):\n",
    "    response = ollama.chat(\n",
    "        model='llama3.1',\n",
    "        messages=[{\"role\": \"system\", \"content\": f\"today's date is {today()}\"}, {'role': 'user', 'content': query}],\n",
    "        tools=tools\n",
    "    )\n",
    "    content = response[\"message\"][\"content\"]\n",
    "    if content:\n",
    "        return content\n",
    "    fn = response[\"message\"][\"tool_calls\"][0]\n",
    "    print(fn)\n",
    "    return call_function(fn)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 .  Give me last year's stock price data.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[*********************100%***********************]  1 of 1 completed"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'function': {'name': 'getHistoricalPrice', 'arguments': {'end': '2023-09-01', 'start': '2022-09-02'}}}\n",
      "Date\n",
      "2022-09-02    155.809998\n",
      "2022-09-06    154.529999\n",
      "2022-09-07    155.960007\n",
      "2022-09-08    154.460007\n",
      "2022-09-09    157.369995\n",
      "                 ...    \n",
      "2023-08-25    178.610001\n",
      "2023-08-28    180.190002\n",
      "2023-08-29    184.119995\n",
      "2023-08-30    187.649994\n",
      "2023-08-31    187.869995\n",
      "Name: Close, Length: 250, dtype: float64\n",
      "---\n",
      "2 .  Lowest price in 2023\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "[*********************100%***********************]  1 of 1 completed"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'function': {'name': 'computeStatistic', 'arguments': {'end': '2023-12-31', 'start': '2023-01-01', 'statistic': 'min'}}}\n",
      "125.0199966430664\n",
      "---\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "queries = [\n",
    "    \"Give me last year's stock price data.\",\n",
    "    \"Lowest price in 2023\",\n",
    "]\n",
    "for i,q in enumerate(queries):\n",
    "    print(i+1, \". \", q)\n",
    "    print(generate_response(q))\n",
    "    print(\"---\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env-dspy",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}