chatbot-stock-buddy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gbeJd3rOc6I1"
   },
   "source": [
    "## Libraries + Env Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A7MmmWa0BTXq"
   },
   "outputs": [],
   "source": [
    "!pip install langchain unstructured openai faiss-cpu tiktoken comet comet_llm yfinance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fAZ11rRwyFnc"
   },
   "outputs": [],
   "source": [
    "import seaborn as sns            # plots\n",
    "import matplotlib.pyplot as plt   # plots\n",
    "import yfinance as yf         # import Yahoo! Finance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vbp2qyC_t3r5"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import unstructured\n",
    "# import streamlit as st\n",
    "# from streamlit_chat import message\n",
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "from langchain.chains import ConversationalRetrievalChain\n",
    "from langchain.document_loaders.csv_loader import CSVLoader\n",
    "from langchain.document_loaders.excel import UnstructuredExcelLoader\n",
    "from langchain.vectorstores import FAISS\n",
    "from langchain.agents import AgentType, initialize_agent\n",
    "from langchain.chat_models import ChatOpenAI\n",
    "import tempfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1WM1b5k-CrjK"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"<enter key>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WwzzjT7Uxqvx"
   },
   "source": [
    "## Create Functions for Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Wvzazsoc7o8W"
   },
   "outputs": [],
   "source": [
    "class SecurityBuddy:\n",
    "\n",
    "  def __init__(self, tickers, market_comparison=False):\n",
    "    self.tickers = tickers\n",
    "\n",
    "    if market_comparison==True:\n",
    "      self.market_indices = ['^DJI', '^IXIC', '^GSPC'] # Dow Jones, Nasdaq and S&P500\n",
    "    else:\n",
    "      self.market_indices = None\n",
    "\n",
    "    self.symbols = self.get_symbols()\n",
    "\n",
    "  def get_symbols(self):\n",
    "    symbols = []\n",
    "    if isinstance(self.tickers, list):\n",
    "      tickers = [x.upper() for x in self.tickers]\n",
    "      tickers.sort()\n",
    "    else:\n",
    "      tickers = [self.tickers.upper()]\n",
    "\n",
    "    if self.market_indices is not None:\n",
    "        for index in self.market_indices:\n",
    "          self.tickers.append(index)\n",
    "        symbols = self.tickers\n",
    "    else:\n",
    "        symbols = self.tickers\n",
    "    return symbols\n",
    "\n",
    "  def get_security_financial_information(self):\n",
    "\n",
    "    # parameters to download\n",
    "    fin_info = [\"shortName\", \"sector\", \"industry\", \"quoteType\", \"exchange\", \"totalAssets\", \"marketCap\", \"beta\", \"trailingPE\", \"volume\", \"averageVolume\", \"fiftyTwoWeekLow\", \"fiftyTwoWeekHigh\", \"dividendRate\", \"phone\"]\n",
    "\n",
    "    #  create dataframe\n",
    "    yf_info = pd.DataFrame(index = fin_info, columns = self.symbols)\n",
    "\n",
    "    for i in self.symbols:\n",
    "      l = []             # initialize\n",
    "      x = yf.Ticker(i)   # get ticker info\n",
    "      for j in fin_info:\n",
    "        if 'date' in j.lower():\n",
    "          d = pd.to_datetime(x.info[j])\n",
    "          if d is not None:\n",
    "            l.append(d.strftime(\"%Y-%m-%d\"))  # format date\n",
    "        else:\n",
    "          try:      # some parameters error\n",
    "            l.append(x.info[j])\n",
    "          except:   # ignore error and continue\n",
    "            l.append(\"\")\n",
    "\n",
    "      yf_info[i] = l\n",
    "\n",
    "    return yf_info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i85gY7SqSbZ3"
   },
   "outputs": [],
   "source": [
    "## test function\n",
    "# tickers = ['MSFT', 'GOOGL']\n",
    "# buddy = SecurityBuddy(tickers, False)\n",
    "# info = buddy.get_security_financial_information()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Qy6aEWQu_IBz"
   },
   "outputs": [],
   "source": [
    "#  SET ticker, yf_period and yf_interval\n",
    "#  ticker - securities to download\n",
    "# ticker = [\"AAPL\", \"GOOG\"]\n",
    "\n",
    "\n",
    "#  time period and interval\n",
    "# yf_period = \"5y\"   # 1d,5d,1mo,3mo,6mo,1y,2y,5y,10y,ytd,max\n",
    "# yf_interval = \"1m\"    # 1m,2m,5m,15m,30m,60m,90m,1h,1d,5d,1wk,1mo,3mo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rvxngVlcA5d9"
   },
   "source": [
    "## Create Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BVt0_uzsA_Dz"
   },
   "outputs": [],
   "source": [
    "from typing import Type\n",
    "\n",
    "from langchain.tools import BaseTool\n",
    "from pydantic import BaseModel, Field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iIf1IpyaA80z"
   },
   "outputs": [],
   "source": [
    "class FinancialBuddyInput(BaseModel):\n",
    "    \"\"\"Inputs for tool class\"\"\"\n",
    "\n",
    "    tickers: Union[str, list] = Field(description=\"Ticker symbol(s) of the stock(s)\")\n",
    "    market_comparison: bool = Field(description=\"boolean value for whether or not to get market information too\")\n",
    "\n",
    "\n",
    "class CurrentFinancialInformationTool(BaseTool):\n",
    "    name = \"get_security_financial_information\"\n",
    "    description = \"\"\"\n",
    "        Useful when you want to get high-level financial information about a stock.\n",
    "        You should enter the stock/index ticker symbols recognized by yahoo finance\n",
    "        \"\"\"\n",
    "    args_schema: Type[BaseModel] = FinancialBuddyInput\n",
    "\n",
    "    def _run(self, tickers: Union[str, list], market_comparison: bool):\n",
    "        buddy = SecurityBuddy(tickers, market_comparison)\n",
    "        response = buddy.get_security_financial_information()\n",
    "        return response\n",
    "\n",
    "    def _arun(self, ticker: str):\n",
    "        raise NotImplementedError(\"get_security_financial_information does not support async\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LQfpy2yrDJsT"
   },
   "source": [
    "## Create Default Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c0USrPdnDPMw"
   },
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo-16k\", temperature=0.9)\n",
    "\n",
    "tools = [CurrentFinancialInformationTool()]\n",
    "\n",
    "agent = initialize_agent(tools, llm, agent=AgentType.OPENAI_FUNCTIONS, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8KLRg10eDibP",
    "outputId": "23508872-d8c8-41ed-fd0e-d112cf294f88"
   },
   "outputs": [],
   "source": [
    "response = agent.run(\n",
    "    \"Which company has a higher market cap? Microsoft or Google?\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "djX4WGOOWbgq",
    "outputId": "13ca520f-c013-4964-a47e-ffd3295a3171"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "type": "string"
      },
      "text/plain": [
       "\"Microsoft has a higher market cap compared to Google. Microsoft's market cap is $2,763,704,500,224, while Google's market cap is $1,641,913,122,816.\""
      ]
     },
     "execution_count": 144,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "axrORfWqdCeI"
   },
   "source": [
    "## WIP - Create Custom Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h4T_kmLodYrH"
   },
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor, BaseSingleActionAgent, Tool\n",
    "from typing import Any, List, Tuple, Union\n",
    "\n",
    "from langchain.schema import AgentAction, AgentFinish"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "i1e5LfujdZ-k"
   },
   "outputs": [],
   "source": [
    "class CustomAgent(BaseSingleActionAgent):\n",
    "    \"\"\"Custom Agent.\"\"\"\n",
    "\n",
    "    @property\n",
    "    def input_keys(self):\n",
    "        return [\"input\"]\n",
    "\n",
    "    def plan(\n",
    "        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any\n",
    "    ) -> Union[AgentAction, AgentFinish]:\n",
    "        \"\"\"Given input, decided what to do.\n",
    "\n",
    "        Args:\n",
    "            intermediate_steps: Steps the LLM has taken to date,\n",
    "                along with observations\n",
    "            **kwargs: User inputs.\n",
    "\n",
    "        Returns:\n",
    "            Action specifying what tool to use.\n",
    "        \"\"\"\n",
    "        return AgentAction(tool=tool, tool_input=kwargs[\"input\"], log=\"\")\n",
    "\n",
    "    async def aplan(\n",
    "        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any\n",
    "    ) -> Union[AgentAction, AgentFinish]:\n",
    "        \"\"\"Given input, decided what to do.\n",
    "\n",
    "        Args:\n",
    "            intermediate_steps: Steps the LLM has taken to date,\n",
    "                along with observations\n",
    "            **kwargs: User inputs.\n",
    "\n",
    "        Returns:\n",
    "            Action specifying what tool to use.\n",
    "        \"\"\"\n",
    "        return AgentAction(tool=tool, tool_input=kwargs[\"input\"], log=\"\")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "BN29akpA1--c"
   ],
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}