{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "client = openai.OpenAI(\n",
    "    base_url = \"https://api.fireworks.ai/inference/v1\",\n",
    "    api_key = \"YOUR_API_KEY\"\n",
    ")\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": f\"You are a helpful assistant with access to functions.\" \n",
    "     \t\t\t\t\t\t\t\t\t\t\t\t\t\t\t\"Use them if required.\"},\n",
    "    {\"role\": \"user\", \"content\": \"What are Nike's net income in 2022?\"}\n",
    "]\n",
    "\n",
    "tools = [\n",
    "    {\n",
    "        \"type\": \"function\",\n",
    "        \"function\": {\n",
    "            # name of the function \n",
    "            \"name\": \"get_financial_data\",\n",
    "            # a good, detailed description for what the function is supposed to do\n",
    "            \"description\": \"Get financial data for a company given the metric and year.\",\n",
    "            # a well defined json schema: https://json-schema.org/learn/getting-started-step-by-step#define\n",
    "            \"parameters\": {\n",
    "                # for OpenAI compatibility, we always declare a top level object for the parameters of the function\n",
    "                \"type\": \"object\",\n",
    "                # the properties for the object would be any arguments you want to provide to the function\n",
    "                \"properties\": {\n",
    "                    \"metric\": {\n",
    "                        # JSON Schema supports string, number, integer, object, array, boolean and null\n",
    "                        # for more information, please check out https://json-schema.org/understanding-json-schema/reference/type\n",
    "                        \"type\": \"string\",\n",
    "                        # You can restrict the space of possible values in an JSON Schema\n",
    "                        # you can check out https://json-schema.org/understanding-json-schema/reference/enum for more examples on how enum works\n",
    "                        \"enum\": [\"net_income\", \"revenue\", \"ebdita\"],\n",
    "                    },\n",
    "                    \"financial_year\": {\n",
    "                        \"type\": \"integer\", \n",
    "                        # If the model does not understand how it is supposed to fill the field, a good description goes a long way \n",
    "                        \"description\": \"Year for which we want to get financial data.\"\n",
    "                    },\n",
    "                    \"company\": {\n",
    "                        \"type\": \"string\",\n",
    "                        \"description\": \"Name of the company for which we want to get financial data.\"\n",
    "                    }\n",
    "                },\n",
    "                # You can specify which of the properties from above are required\n",
    "                # for more info on `required` field, please check https://json-schema.org/understanding-json-schema/reference/object#required\n",
    "                \"required\": [\"metric\", \"financial_year\", \"company\"],\n",
    "            },\n",
    "        },\n",
    "    }\n",
    "]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"content\": null,\n",
      "    \"role\": \"assistant\",\n",
      "    \"function_call\": null,\n",
      "    \"tool_calls\": [\n",
      "        {\n",
      "            \"id\": \"call_W4IUqQRFF9vYINQ74tfBwmqr\",\n",
      "            \"function\": {\n",
      "                \"arguments\": \"{\\\"metric\\\": \\\"net_income\\\", \\\"financial_year\\\": 2022, \\\"company\\\": \\\"Nike\\\"}\",\n",
      "                \"name\": \"get_financial_data\"\n",
      "            },\n",
      "            \"type\": \"function\",\n",
      "            \"index\": 0\n",
      "        }\n",
      "    ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "\n",
    "chat_completion = client.chat.completions.create(\n",
    "    model=\"accounts/fireworks/models/firefunction-v2\",\n",
    "    messages=messages,\n",
    "    tools=tools,\n",
    "    temperature=0.1\n",
    ")\n",
    "print(chat_completion.choices[0].message.model_dump_json(indent=4))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}