function_calling/OpenAI_function_calling.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Copyright 2024 Kensho Technologies, LLC"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "-zOQ3ZfKSrT2"
   },
   "source": [
    "# OpenAI Function Calling\n",
    "**_GPT to retrieve data from the LLM-ready API using the kFinance python library!_**\n",
    "\n",
    "What you'll need to run this notebook:\n",
    "\n",
    "1.   kFinance credentials\n",
    "2.   An OpenAI API key\n",
    "\n",
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/kensho-technologies/llm-ready-api-examples/blob/main/function_calling/OpenAI_function_calling.ipynb\"><img src=\"../images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "ruP6MHfzR2Y4"
   },
   "source": [
    "## Install dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "collapsed": true,
    "id": "YAwU0_7TRhV7",
    "jupyter": {
     "outputs_hidden": true
    },
    "outputId": "e6ebf89c-7680-4852-e52b-c99610107cdf"
   },
   "outputs": [],
   "source": [
    "# install the latest version of kFinance package\n",
    "%pip install https://kfinance.kensho.com/static/kensho_finance.tar.gz\n",
    "# install the LLM Python package\n",
    "%pip install openai"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "e_G2p3tTSHeo"
   },
   "source": [
    "# Instantiate kFinance Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7QPnfPuX87rg"
   },
   "outputs": [],
   "source": [
    "# import the kfinance client\n",
    "from kensho_finance.kfinance import Client\n",
    "# import standard libraries\n",
    "import functools\n",
    "import types\n",
    "import json\n",
    "import sys\n",
    "# check if the current environment is a Google Colab\n",
    "try:\n",
    "  import google.colab\n",
    "  IN_GOOGLE_COLAB = True\n",
    "except:\n",
    "  IN_GOOGLE_COLAB = False\n",
    "\n",
    "# initialize the kfinance client with one of the following:\n",
    "# 1. your kensho refresh token\n",
    "# 2. your kensho client id and kensho private key\n",
    "# 3. automated login (not accessible on Google Collab)\n",
    "if IN_GOOGLE_COLAB:\n",
    "    kensho_refresh_token = \"\"\n",
    "    assert kensho_refresh_token != \"\", \"kensho refresh token is empty! Make sure to enter your kensho refresh token above\"\n",
    "    kfinance_client = Client(refresh_token=kensho_refresh_token)\n",
    "\n",
    "    # kensho_client_id = \"\"\n",
    "    # kensho_private_key = \"\"\n",
    "    # assert kensho_client_id != \"\", \"kensho client id is empty! Make sure to enter your kensho client id above\"\n",
    "    # assert kensho_private_key != \"\", \"kensho private key is empty! Make sure to enter your kensho private key above\"\n",
    "    # kfinance_client = Client(client_id=kensho_client_id, private_key=kensho_private_key)\n",
    "else:\n",
    "    kfinance_client = Client()\n",
    "\n",
    "# the prompt given to all LLMs to instruct them how to use tools to call the kFinance API\n",
    "SYSTEM_PROMPT = \"You are an LLM trying to help financial analysts. Use the supplied tools to assist the user. Always use the `get_latest` function when asked about the last or most recent quarter or time period etc. Always use the `get_latest` function when a tool requires a time parameter and the time is unspecified in the question\"\n",
    "# the message shown to users to prompt them to input a message\n",
    "USER_INPUT_PROMPT = \"Enter your message and press the [return] key\\n\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "LvQy_66Jioys"
   },
   "source": [
    "# OpenAI Function Calling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WkoRuipOEQz6"
   },
   "outputs": [],
   "source": [
    "# import OpenAI\n",
    "from openai import OpenAI\n",
    "# create a struct similar to the OpenAI response class, for manually creating responses\n",
    "from collections import namedtuple\n",
    "Message = namedtuple('Message', 'role content tool_calls')\n",
    "\n",
    "# the OpenAIChat class is used to create a chat loop that automatically executes tool calls\n",
    "class OpenAIChat:\n",
    "  def __init__(self, kfinance_client: Client):\n",
    "    # initialize OpenAI with your OpenAI API key\n",
    "    openai_api_key = \"\" # replace with your own key\n",
    "    assert openai_api_key != \"\", \"OpenAI API key is empty! Make sure to enter your OpenAI API key above\"\n",
    "    # good client config for OpenAI. There are other ways to access OpenAI\n",
    "    self.openai = OpenAI(\n",
    "      api_key=openai_api_key,\n",
    "    )\n",
    "    # initialize the kFinance client\n",
    "    self.kfinance_client = kfinance_client\n",
    "    # initialize tools and tool descriptions\n",
    "    self.tools = kfinance_client.tools\n",
    "    self.tool_descriptions = kfinance_client.openai_tool_descriptions\n",
    "    # set the first message to include the system prompt\n",
    "    self.messages = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": SYSTEM_PROMPT,\n",
    "        },\n",
    "    ]\n",
    "\n",
    "  def print_response(self):\n",
    "      \"\"\"Print and return response\"\"\"\n",
    "      # try to send the message history to OpenAI and get the response\n",
    "      try:\n",
    "          response = self.openai.chat.completions.create(\n",
    "              model=\"gpt-4o\", # you can use any OpenAI model that supports function calling\n",
    "              messages=self.messages,\n",
    "              tools=self.tool_descriptions,\n",
    "          )\n",
    "          response_message = response.choices[0].message\n",
    "          self.messages.append(response_message)\n",
    "      # if there's an error, manually create a response with the error\n",
    "      except Exception as e:\n",
    "          self.messages.append({\"role\": \"user\", \"content\": f\"you had an error: {str(e)}\"})\n",
    "          response_message = Message(\"assistant\", f\"an error occured: {str(e)}\", None)\n",
    "      # if the response contains text (and not just tool calls), print out the text\n",
    "      if response_message.content is not None:\n",
    "            sys.stdout.write(response_message.content)\n",
    "      return response_message\n",
    "\n",
    "  def print_responses(self, user_input: str) -> list[str]:\n",
    "      \"\"\"Print responses and call tools\"\"\"\n",
    "      # append the current user input as a message to the message history\n",
    "      self.messages.append({\"role\": \"user\", \"content\": user_input})\n",
    "      # get the OpenAI response to the message history\n",
    "      response_message = self.print_response()\n",
    "      # while the response has tool calls\n",
    "      while response_message.tool_calls is not None:\n",
    "          # for each tool call, execute the function and arguments specified in the tool call\n",
    "          # and append the output as a message to the message history\n",
    "          for tool_call in response_message.tool_calls:\n",
    "              function = tool_call.function.name\n",
    "              arguments = json.loads(tool_call.function.arguments)\n",
    "              # try to execute the function and arguments specified in the tool call\n",
    "              try:\n",
    "                  output = self.tools[function](**arguments)\n",
    "                  # append the output as a message to the message history\n",
    "                  self.messages.append(\n",
    "                      {\n",
    "                          \"role\": \"tool\",\n",
    "                          \"content\": json.dumps({\"output\": str(output)}),\n",
    "                          \"tool_call_id\": tool_call.id,\n",
    "                      }\n",
    "                  )\n",
    "              # if there's an exception thrown while executing the function,\n",
    "              # append the exception as a message to the message history\n",
    "              except Exception as e:\n",
    "                  self.messages.append(\n",
    "                      {\n",
    "                          \"role\": \"tool\",\n",
    "                          \"content\": json.dumps({\"output\": str(e)}),\n",
    "                          \"tool_call_id\": tool_call.id,\n",
    "                      }\n",
    "                  )\n",
    "          # get a new response\n",
    "          response_message = self.print_response()\n",
    "      return None\n",
    "\n",
    "  def start_chatting(self) -> None:\n",
    "      \"\"\"Open chat shell\"\"\"\n",
    "      # prompt for user input and get the OpenAI response in a loop\n",
    "      while True:\n",
    "          user_input = input(USER_INPUT_PROMPT)\n",
    "          self.print_responses(user_input)\n",
    "          sys.stdout.write(\"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 510
    },
    "id": "t2x5XdwyKYj-",
    "outputId": "ec50e672-55cf-40a8-f96d-2b654e60538a"
   },
   "outputs": [],
   "source": [
    "# instantiate the OpenAIChat with the kfinance client\n",
    "openai_chat = OpenAIChat(kfinance_client)\n",
    "# start chatting with the OpenAIChat\n",
    "openai_chat.start_chatting()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}