{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import OpenAI key with helper function\n",
    "from helper import get_openai_api_key\n",
    "\n",
    "OPENAI_API_KEY = get_openai_api_key()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#A lot of modules use async and we want them to be compatible with Jupyter notebook\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Ehlers-Danlos-1\\\\2024_EDS_1.pdf', 'Ehlers-Danlos-1\\\\2024_EDS_2.pdf', 'Ehlers-Danlos-1\\\\2024_EDS_3.pdf', 'Ehlers-Danlos-1\\\\2024_EDS_4.pdf', 'Ehlers-Danlos-1\\\\2024_EDS_5.pdf', 'Ehlers-Danlos-1\\\\2024_EDS_6.pdf', 'Ehlers-Danlos-1\\\\2024_EDS_7.pdf', 'Ehlers-Danlos-1\\\\Unknown_EDS_1.pdf', 'Ehlers-Danlos-1\\\\Unknown_EDS_2.pdf', 'Ehlers-Danlos-1\\\\Unknown_EDS_3.pdf', 'Ehlers-Danlos-1\\\\Unknown_EDS_4.pdf', 'Ehlers-Danlos-1\\\\Unknown_EDS_5.pdf']\n",
      "['2024_EDS_1.pdf', '2024_EDS_2.pdf', '2024_EDS_3.pdf', '2024_EDS_4.pdf', '2024_EDS_5.pdf', '2024_EDS_6.pdf', '2024_EDS_7.pdf', 'Unknown_EDS_1.pdf', 'Unknown_EDS_2.pdf', 'Unknown_EDS_3.pdf', 'Unknown_EDS_4.pdf', 'Unknown_EDS_5.pdf']\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import glob\n",
    "\n",
    "# Define the path to the directory containing the PDF files\n",
    "folder_path = 'Ehlers-Danlos-1'\n",
    "\n",
    "# Get the list of all PDF files in the directory\n",
    "pdf_files = glob.glob(os.path.join(folder_path, '*.pdf'))\n",
    "print(pdf_files)\n",
    "\n",
    "# Extract just the filenames (optional)\n",
    "pdf_filenames = [os.path.basename(pdf) for pdf in pdf_files]\n",
    "\n",
    "# Print the list of PDF filenames\n",
    "print(pdf_filenames)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Getting tools for paper: Ehlers-Danlos-1\\2024_EDS_1.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\2024_EDS_2.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\2024_EDS_3.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\2024_EDS_4.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\2024_EDS_5.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\2024_EDS_6.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\2024_EDS_7.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\Unknown_EDS_1.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\Unknown_EDS_2.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\Unknown_EDS_3.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\Unknown_EDS_4.pdf\n",
      "Getting tools for paper: Ehlers-Danlos-1\\Unknown_EDS_5.pdf\n"
     ]
    }
   ],
   "source": [
    "from utils import get_doc_tools\n",
    "from pathlib import Path\n",
    "\n",
    "# Ensure function names are within the allowed length limit\n",
    "def truncate_function_name(name, max_length=64):\n",
    "    return name if len(name) <= max_length else name[:max_length]\n",
    "\n",
    "paper_to_tools_dict = {}\n",
    "for pdf in pdf_files:\n",
    "    print(f\"Getting tools for paper: {pdf}\")\n",
    "    vector_tool, summary_tool = get_doc_tools(pdf, Path(pdf).stem)\n",
    "    #vector_tool, summary_tool = get_doc_tools(pdf, truncate_function_name(Path(pdf).stem))\n",
    "    paper_to_tools_dict[pdf] = [vector_tool, summary_tool]\n",
    "    #print(vector_tool)\n",
    "    #print(summary_tool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_tools = [t for pdf in pdf_files for t in paper_to_tools_dict[pdf]]\n",
    "#all_tools = [truncate_function_name(tool) for tool in all_tools]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define an \"object\" index and retriever over these tools\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core.objects import ObjectIndex\n",
    "\n",
    "obj_index = ObjectIndex.from_objects(\n",
    "    all_tools,\n",
    "    index_cls=VectorStoreIndex,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "obj_retriever = obj_index.as_retriever(similarity_top_k=3)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent import FunctionCallingAgentWorker\n",
    "from llama_index.core.agent import AgentRunner\n",
    "\n",
    "agent_worker = FunctionCallingAgentWorker.from_tools(\n",
    "    tool_retriever=obj_retriever,\n",
    "    llm=llm, \n",
    "    verbose=True\n",
    ")\n",
    "agent = AgentRunner(agent_worker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'agent' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[1], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m response \u001b[38;5;241m=\u001b[39m \u001b[43magent\u001b[49m\u001b[38;5;241m.\u001b[39mquery(\n\u001b[0;32m      2\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDo people with EDS suffer from dislocations, and if so, how do they manifest?\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m      3\u001b[0m )\n\u001b[0;32m      4\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;28mstr\u001b[39m(response))\n",
      "\u001b[1;31mNameError\u001b[0m: name 'agent' is not defined"
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n",
      "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n",
      "\u001b[1;31mClick <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. \n",
      "\u001b[1;31mView Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "\n",
    "response = agent.query(\n",
    "    \"Do people with EDS suffer from dislocations, and if so, how do they manifest?\"\n",
    ")\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}