Cases.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Case I: Protein design and analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "PROJECT_ROOT = os.getcwd()\n",
    "load_dotenv(os.path.join(PROJECT_ROOT, 'example.env'), override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hjj/miniconda3/envs/toolsKG/lib/python3.9/site-packages/Bio/pairwise2.py:278: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from scripts.custom_kg_retrievers import KGTableRetriever\n",
    "from scripts.retrieve_tool_info import execute_tool, check_toxicity\n",
    "from scripts.utils import load_index_from_storage_dir, construct_question\n",
    "from app.llms.planning import generate_tools_plan, generate_plan_input, generate_output\n",
    "from app.llms.extraction import extract_tool_path, extract_answer_parts\n",
    "from app.tools.safety.toxicity_checker import check_toxicity_by_type\n",
    "from app.tools.safety import extract_smiles_or_sequences\n",
    "from app.tools import run_tool\n",
    "from app.core.config import Config\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Settings\n",
    "from llama_index.core import QueryBundle\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "config = Config()\n",
    "Settings.embed_model = OpenAIEmbedding(model=config.EMBEDDING_MODEL_NAME)\n",
    "Settings.llm = OpenAI(temperature=0, model=config.MODEL_NAME)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.core.indices.loading:Loading all indices.\n",
      "Knowledge Graph Index loaded successfully from KG/storage_graph_large\n"
     ]
    }
   ],
   "source": [
    "question = \"Design an alpha-structure protein consisting of 200 amino acids and generate its corresponding sequence. Then, convert this sequence into a 3D structure, producing a PDB file and calculating the residue accuracy (e.g., pLDDT values). Analyze the secondary structure composition of the protein using the generated PDB file. Next, calculate both the unfolding force and energy from the protein sequence. Finally, Calculates the Anisotropic Network Model to analyze the first 10 vibrational frequencies and dynamic motions of the protein.\"\n",
    "is_add_filename = False\n",
    "file_path_list = []\n",
    "\n",
    "question = construct_question(question, file_path_list, is_add_filename)\n",
    "storage_context, kg_index = load_index_from_storage_dir(config.PERSIST_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "['SequenceToPdb', 'ProteinSolubilityPredictor', 'PPI3DTwo', 'CalculateEnergyFromSequence', 'DesignProteinAlphaBeta', 'DesignProteinAlpha', 'ProteinMotifAnalysis', 'CalculateProteinANM', 'CalculateForceEnergyFromSequence', 'AnalyzeProteinStructure']\n"
     ]
    }
   ],
   "source": [
    "kg_table_retriever = KGTableRetriever(\n",
    "    index=kg_index, \n",
    "    similarity_top_k=6, \n",
    "    graph_store_query_depth=3, \n",
    "    max_tools=10\n",
    ")        \n",
    "query_bundle = QueryBundle(query_str=question)\n",
    "tools_retrieved = await kg_table_retriever._retrieve(query_bundle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "['DesignProteinAlpha', 'SequenceToPdb', 'AnalyzeProteinStructure', 'CalculateForceEnergyFromSequence', 'CalculateProteinANM']\n"
     ]
    }
   ],
   "source": [
    "responses_content = \"\"\n",
    "total_output = \"\"\n",
    "\n",
    "tools_plan = await generate_tools_plan(\n",
    "    question=question, \n",
    "    tools_retrieved_list=tools_retrieved, \n",
    "    model=config.MODEL_NAME, \n",
    "    previous_answer=total_output, \n",
    ")\n",
    "tools_path = extract_tool_path(tools_plan)\n",
    "print(tools_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "responses_content = \"\"\n",
    "total_output = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing tool: DesignProteinAlpha\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: 200\n",
      "current_output: \n",
      "\n",
      "Now we have run the **DesignProteinAlpha** and obtained result\n",
      " ## Protein Design Result\n",
      "- **Protein Filename**: `alpha_protein`\n",
      "- **CATH Annotation**: `alpha`\n",
      "- **Sequence Length**: `200`\n",
      "\n",
      "**Protein Sequence**:  \n",
      "`SEEEKIKEIREKLKDAVLEAYELFRQGRPQEGEQRLIEAMKEAAKELNVKLSDEEIQREATRLLVEAFREQAKKEGIDEDLAVIAFAELMRALGDTSPLAQEVLREAEELRQKDVLGNKIKEIIEMAKKGNYDEAEKKFKQLKETVDKLIEELRRAADLYASRVSPEEAAAFRQRIDTLLSEMKEKVEKAENEIKKIKEK`.\n",
      "Executing tool: SequenceToPdb\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: SEEEKIKEIREKLKDAVLEAYELFRQGRPQEGEQRLIEAMKEAAKELNVKLSDEEIQREATRLLVEAFREQAKKEGIDEDLAVIAFAELMRALGDTSPLAQEVLREAEELRQKDVLGNKIKEIIEMAKKGNYDEAEKKFKQLKETVDKLIEELRRAADLYASRVSPEEAAAFRQRIDTLLSEMKEKVEKAENEIKKIKEK\n",
      "current_output: \n",
      "\n",
      "Now we have run the **SequenceToPdb** and obtained result\n",
      " The file has been stored with the name alpha_protein.pdb.\n",
      "The mutated plddt values were 78.67665230769232.\n",
      "Executing tool: AnalyzeProteinStructure\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: alpha_protein.pdb\n",
      "current_output: \n",
      "\n",
      "Now we have run the **AnalyzeProteinStructure** and obtained result\n",
      " {\n",
      "    \"H\": 85.5,\n",
      "    \"B\": 0.0,\n",
      "    \"E\": 0.0,\n",
      "    \"G\": 0.0,\n",
      "    \"I\": 0.0,\n",
      "    \"T\": 7.000000000000001,\n",
      "    \"S\": 1.5,\n",
      "    \"P\": 0.0,\n",
      "    \"-\": 6.0\n",
      "}.\n",
      "Executing tool: CalculateForceEnergyFromSequence\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: SEEEKIKEIREKLKDAVLEAYELFRQGRPQEGEQRLIEAMKEAAKELNVKLSDEEIQREATRLLVEAFREQAKKEGIDEDLAVIAFAELMRALGDTSPLAQEVLREAEELRQKDVLGNKIKEIIEMAKKGNYDEAEKKFKQLKETVDKLIEELRRAADLYASRVSPEEAAAFRQRIDTLLSEMKEKVEKAENEIKKIKEK\n",
      "current_output: \n",
      "\n",
      "Now we have run the **CalculateForceEnergyFromSequence** and obtained result\n",
      " ### Force and Energy Calculation Result\n",
      "\n",
      "- **Sequence**: `SEEEKIKEIREKLKDAVLEAYELFRQGRPQEGEQRLIEAMKEAAKELNVKLSDEEIQREATRLLVEAFREQAKKEGIDEDLAVIAFAELMRALGDTSPLAQEVLREAEELRQKDVLGNKIKEIIEMAKKGNYDEAEKKFKQLKETVDKLIEELRRAADLYASRVSPEEAAAFRQRIDTLLSEMKEKVEKAENEIKKIKEK`\n",
      "\n",
      "- **Unfolding Force**: `0.379`\n",
      "- **Energy**: `0.371`.\n",
      "Executing tool: CalculateProteinANM\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: alpha_protein\n",
      "current_output: \n",
      "\n",
      "Now we have run the **CalculateProteinANM** and obtained result\n",
      " ### ANM Calculation Results for `alpha_protein.pdb`\n",
      "\n",
      "- **Number of Modes**: `10`\n",
      "- **Cutoff Distance**: `12.0` Å\n",
      "\n",
      "**Eigenvalues**:\n",
      "\n",
      "- Mode 1: `0.8081`\n",
      "- Mode 2: `0.9031`\n",
      "- Mode 3: `1.1123`\n",
      "- Mode 4: `1.6971`\n",
      "- Mode 5: `2.0849`\n",
      "- Mode 6: `2.4023`\n",
      "- Mode 7: `3.2085`\n",
      "- Mode 8: `3.6356`\n",
      "- Mode 9: `3.9292`\n",
      "- Mode 10: `4.6979`\n",
      ".\n"
     ]
    }
   ],
   "source": [
    "for tool_name in tools_path:\n",
    "    print(f\"Executing tool: {tool_name}\")\n",
    "    current_output = await execute_tool(\n",
    "        tool_name, \n",
    "        question, \n",
    "        question, \n",
    "        file_path_list, \n",
    "        config.MODEL_NAME, \n",
    "        kg_table_retriever, \n",
    "        total_output\n",
    "        )\n",
    "    print(f\"current_output: {current_output}\")\n",
    "    total_output += current_output\n",
    "    responses_content += current_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    }
   ],
   "source": [
    "answer_text = await generate_output(question, total_output, model=config.MODEL_NAME)\n",
    "is_completed, current_final_answer = extract_answer_parts(answer_text)\n",
    "responses_content += f\"\\nTask completion status: {is_completed}\\nFinal Answer: {current_final_answer}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed\n",
      "Yes, your problem has been solved.\n",
      "\n",
      "- **Protein Design Result**:\n",
      "  - **Protein Filename**: `alpha_protein`\n",
      "  - **CATH Annotation**: `alpha`\n",
      "  - **Sequence Length**: `200`\n",
      "  - **Protein Sequence**: `SEEEKIKEIREKLKDAVLEAYELFRQGRPQEGEQRLIEAMKEAAKELNVKLSDEEIQREATRLLVEAFREQAKKEGIDEDLAVIAFAELMRALGDTSPLAQEVLREAEELRQKDVLGNKIKEIIEMAKKGNYDEAEKKFKQLKETVDKLIEELRRAADLYASRVSPEEAAAFRQRIDTLLSEMKEKVEKAENEIKKIKEK`\n",
      "\n",
      "- **PDB File**:\n",
      "  - The file has been stored with the name `alpha_protein.pdb`\n",
      "  - **plddt Value**: `78.67665230769232`\n",
      "\n",
      "- **Secondary Structure Composition**:\n",
      "  - **H (Helix)**: `85.5%`\n",
      "  - **B (Isolated beta-bridge)**: `0.0%`\n",
      "  - **E (Extended strand)**: `0.0%`\n",
      "  - **G (3-helix (310 helix))**: `0.0%`\n",
      "  - **I (5-helix (pi helix))**: `0.0%`\n",
      "  - **T (Turn)**: `7.0%`\n",
      "  - **S (Bend)**: `1.5%`\n",
      "  - **P (Polyproline II helix)**: `0.0%`\n",
      "  - **- (Coil)**: `6.0%`\n",
      "\n",
      "- **Unfolding Force and Energy**:\n",
      "  - **Unfolding Force**: `0.379`\n",
      "  - **Energy**: `0.371`\n",
      "\n",
      "- **ANM Calculation Results**:\n",
      "  - **Number of Modes**: `10`\n",
      "  - **Cutoff Distance**: `12.0` Å\n",
      "  - **Eigenvalues**:\n",
      "    - Mode 1: `0.8081`\n",
      "    - Mode 2: `0.9031`\n",
      "    - Mode 3: `1.1123`\n",
      "    - Mode 4: `1.6971`\n",
      "    - Mode 5: `2.0849`\n",
      "    - Mode 6: `2.4023`\n",
      "    - Mode 7: `3.2085`\n",
      "    - Mode 8: `3.6356`\n",
      "    - Mode 9: `3.9292`\n",
      "    - Mode 10: `4.6979`\n"
     ]
    }
   ],
   "source": [
    "print(is_completed)\n",
    "print(current_final_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Case II: Chemical reactivity prediction and analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PROJECT_ROOT: /home/hjj/project/SciToolAgent\n",
      "PERSIST_DIR: /home/hjj/project/SciToolAgent/KG/storage_graph_large\n",
      "DATA_FILE_PATH: /home/hjj/project/SciToolAgent/data/toolsKG_full.xlsx\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "PROJECT_ROOT = os.getcwd()\n",
    "\n",
    "load_dotenv(os.path.join(PROJECT_ROOT, 'example.env'), override=True)\n",
    "\n",
    "PERSIST_DIR = os.path.join(PROJECT_ROOT, os.getenv(\"PERSIST_DIR\"))\n",
    "DATA_FILE_PATH = os.path.join(PROJECT_ROOT, os.getenv(\"DATA_FILE_PATH\"))\n",
    "\n",
    "print(\"PROJECT_ROOT:\", PROJECT_ROOT)\n",
    "print(\"PERSIST_DIR:\", PERSIST_DIR)\n",
    "print(\"DATA_FILE_PATH:\", DATA_FILE_PATH)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hjj/miniconda3/envs/toolsKG/lib/python3.9/site-packages/Bio/pairwise2.py:278: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "import logging\n",
    "import os\n",
    "import pdb\n",
    "import sys\n",
    "\n",
    "from scripts.custom_kg_retrievers import KGTableRetriever\n",
    "from scripts.retrieve_tool_info import execute_tool, check_toxicity\n",
    "from scripts.utils import load_index_from_storage_dir, construct_question\n",
    "from app.llms.planning import generate_tools_plan, generate_plan_input, generate_output\n",
    "from app.llms.extraction import extract_tool_path, extract_answer_parts\n",
    "from app.tools.safety.toxicity_checker import check_toxicity_by_type\n",
    "from app.tools.safety import extract_smiles_or_sequences\n",
    "from app.tools import run_tool\n",
    "from app.core.config import Config\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Settings\n",
    "from llama_index.core import QueryBundle\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "# print()\n",
    "config = Config()\n",
    "Settings.embed_model = OpenAIEmbedding(model=config.EMBEDDING_MODEL_NAME)\n",
    "Settings.llm = OpenAI(temperature=0, model=config.MODEL_NAME)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KG/storage_graph_large\n",
      "INFO:llama_index.core.indices.loading:Loading all indices.\n",
      "Knowledge Graph Index loaded successfully from KG/storage_graph_large\n"
     ]
    }
   ],
   "source": [
    "question = \"Please predict the reactivity of chemical reactants, provided in SMILES format in Test.csv. You can try different molecular features, such as RDK fingerprints, Morgan fingerprints and electrical descriptors, and use a machine learning algorithm (such as MLP) to predict reactivity. Tell me which feature yields the best predictive performance.\"\n",
    "\n",
    "is_add_filename = True\n",
    "file_path_list = ['/home/hjj/project/ToolsKG/TempFiles/ChemLab/demo_test.csv']\n",
    "\n",
    "question = construct_question(question, file_path_list, is_add_filename)\n",
    "print(config.PERSIST_DIR)\n",
    "storage_context, kg_index = load_index_from_storage_dir(config.PERSIST_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "['GenerateRDKFingerprintsFromCSV', 'GenerateElectricalDescriptorsFromCSV', 'MLPClassifier', 'RandomForestClassifier', 'GenerateMorganfingerprintsFromCSV', 'RXNPredict', 'AdaBoostClassifier', 'GetFraggleSimilarity', 'MolSimilarity', 'GetSubstructMatches']\n"
     ]
    }
   ],
   "source": [
    "responses_content = \"\"\n",
    "total_output = \"\"\n",
    "\n",
    "kg_table_retriever = KGTableRetriever(\n",
    "    index=kg_index, \n",
    "    similarity_top_k=5, \n",
    "    graph_store_query_depth=3, \n",
    "    max_tools=10\n",
    ")        \n",
    "query_bundle = QueryBundle(query_str=question)\n",
    "tools_retrieved = await kg_table_retriever._retrieve(query_bundle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "['GenerateRDKFingerprintsFromCSV', 'GenerateMorganfingerprintsFromCSV', 'GenerateElectricalDescriptorsFromCSV', 'MLPClassifier', 'MLPClassifier', 'MLPClassifier']\n"
     ]
    }
   ],
   "source": [
    "tools_plan = await generate_tools_plan(\n",
    "    question=question, \n",
    "    tools_retrieved_list=tools_retrieved, \n",
    "    model=config.MODEL_NAME, \n",
    "    previous_answer=total_output, \n",
    ")\n",
    "tools_path = extract_tool_path(tools_plan)\n",
    "print(tools_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing tool: GenerateRDKFingerprintsFromCSV\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: demo_test.csv\n",
      "\n",
      "current_output: Now we have run the GenerateRDKFingerprintsFromCSV and obtained result \n",
      "smiles in the demo_test.csv file have been processed, and the fingerprints features are saved to the path ToolsKG/TempFiles/ChemLab/demo_test_RDK_fingerprints.csv\n",
      ".\n",
      "\n",
      "Executing tool: GenerateMorganfingerprintsFromCSV\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: demo_test.csv\n",
      "\n",
      "current_output: Now we have run the GenerateMorganfingerprintsFromCSV and obtained result \n",
      "SMILES strings in the demo_test.csv file have been processed, and the morgan fingerprints features are saved to the path ToolsKG/TempFiles/ChemLab/demo_test_Morgan_fingerprints.csv\n",
      ".\n",
      "\n",
      "Executing tool: GenerateElectricalDescriptorsFromCSV\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: demo_test.csv\n",
      "\n",
      "current_output: Now we have run the GenerateElectricalDescriptorsFromCSV and obtained result Processed SMILES strings and saved electrical descriptors to ToolsKG/TempFiles/ChemLab/demo_test_electrical_descriptors.csv.\n",
      "\n",
      "Executing tool: MLPClassifier\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: ToolsKG/TempFiles/ChemLab/demo_test_RDK_fingerprints.csv\n",
      "\n",
      "current_output: Now we have run the MLPClassifier and obtained result \n",
      "### MLP Classifier Results for Rdk Fingerprints\n",
      "\n",
      "- **Training Accuracy**: 1.0000\n",
      "- **Test Accuracy**: 0.6818\n",
      "\n",
      ".\n",
      "\n",
      "Executing tool: MLPClassifier\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: ToolsKG/TempFiles/ChemLab/demo_test_Morgan_fingerprints.csv\n",
      "\n",
      "current_output: Now we have run the MLPClassifier and obtained result \n",
      "### MLP Classifier Results for Morgan Fingerprints\n",
      "\n",
      "- **Training Accuracy**: 1.0000\n",
      "- **Test Accuracy**: 0.7727\n",
      "\n",
      ".\n",
      "\n",
      "Executing tool: MLPClassifier\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: ToolsKG/TempFiles/ChemLab/demo_test_electrical_descriptors.csv\n",
      "\n",
      "current_output: Now we have run the MLPClassifier and obtained result \n",
      "### MLP Classifier Results for Electrical Descriptors\n",
      "\n",
      "- **Training Accuracy**: 0.8158\n",
      "- **Test Accuracy**: 0.9091\n",
      "\n",
      ".\n"
     ]
    }
   ],
   "source": [
    "for tool_name in tools_path:\n",
    "    print(f\"Executing tool: {tool_name}\")\n",
    "    current_output = await execute_tool(\n",
    "        tool_name, \n",
    "        question, \n",
    "        question, \n",
    "        file_path_list, \n",
    "        config.MODEL_NAME, \n",
    "        kg_table_retriever, \n",
    "        total_output\n",
    "        )\n",
    "    print(f\"\\ncurrent_output: {current_output}\\n\")\n",
    "    total_output += current_output\n",
    "    responses_content += current_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Completed'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "answer_text = await generate_output(question, total_output, model=config.MODEL_NAME)\n",
    "is_completed, current_final_answer = extract_answer_parts(answer_text)\n",
    "responses_content += f\"\\nTask completion status: {is_completed}\\nFinal Answer: {current_final_answer}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed\n",
      "The problem has been solved. The best predictive performance was achieved using electrical descriptors, with a test accuracy of 0.9091.\n"
     ]
    }
   ],
   "source": [
    "print(is_completed)\n",
    "print(current_final_answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "question = \"Based on the results that electrical descriptors perform best for predicting reactivity, I would like to compare the performance of different machine learning algorithms. Please use multiple algorithms (MLP, AdaBoost, and Random Forest) to predict reactivity and identify which algorithm yields the best accuracy.\"\n",
    "is_add_filename = True\n",
    "file_path_list = ['/home/hjj/project/ToolsKG/TempFiles/ChemLab/demo_test.csv']\n",
    "\n",
    "question = construct_question(question, file_path_list, is_add_filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "['MLPClassifier', 'AdaBoostClassifier', 'RandomForestClassifier', 'GenerateElectricalDescriptorsFromCSV', 'PredictAdsorption', 'PredictStability', 'MOFToSMILES', 'GenerateRDKFingerprintsFromCSV', 'GetFloatingSolventMolecules', 'PlotAdjacencyMatrix']\n"
     ]
    }
   ],
   "source": [
    "responses_content = \"\"\n",
    "total_output = \"\"\n",
    "\n",
    "kg_table_retriever = KGTableRetriever(\n",
    "    index=kg_index, \n",
    "    similarity_top_k=5, \n",
    "    graph_store_query_depth=3, \n",
    "    max_tools=10\n",
    ")        \n",
    "query_bundle = QueryBundle(query_str=question)\n",
    "tools_retrieved = await kg_table_retriever._retrieve(query_bundle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "['GenerateElectricalDescriptorsFromCSV', 'MLPClassifier', 'AdaBoostClassifier', 'RandomForestClassifier']\n"
     ]
    }
   ],
   "source": [
    "tools_plan = await generate_tools_plan(\n",
    "    question=question, \n",
    "    tools_retrieved_list=tools_retrieved, \n",
    "    model=config.MODEL_NAME, \n",
    "    previous_answer=total_output, \n",
    ")\n",
    "tools_path = extract_tool_path(tools_plan)\n",
    "print(tools_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing tool: GenerateElectricalDescriptorsFromCSV\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: demo_test.csv\n",
      "\n",
      "current_output: Now we have run the GenerateElectricalDescriptorsFromCSV and obtained result Processed SMILES strings and saved electrical descriptors to ToolsKG/TempFiles/ChemLab/demo_test_electrical_descriptors.csv.\n",
      "\n",
      "Executing tool: MLPClassifier\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: ToolsKG/TempFiles/ChemLab/demo_test_electrical_descriptors.csv\n",
      "\n",
      "current_output: Now we have run the MLPClassifier and obtained result \n",
      "### MLP Classifier Results for Electrical Descriptors\n",
      "\n",
      "- **Training Accuracy**: 0.7703\n",
      "- **Test Accuracy**: 0.9545\n",
      "\n",
      ".\n",
      "\n",
      "Executing tool: AdaBoostClassifier\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: ToolsKG/TempFiles/ChemLab/demo_test_electrical_descriptors.csv\n",
      "\n",
      "current_output: Now we have run the AdaBoostClassifier and obtained result \n",
      "### AdaBoost Classifier Results for Electrical Descriptors\n",
      "\n",
      "- **Training Accuracy**: 0.9091\n",
      "- **Test Accuracy**: 0.8636\n",
      "\n",
      ".\n",
      "\n",
      "Executing tool: RandomForestClassifier\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: ToolsKG/TempFiles/ChemLab/demo_test_electrical_descriptors.csv\n",
      "\n",
      "current_output: Now we have run the RandomForestClassifier and obtained result \n",
      "### Random Forest Classifier Results for Electrical Descriptors\n",
      "\n",
      "- **Training Accuracy**: 1.0000\n",
      "- **Test Accuracy**: 0.9545\n",
      "\n",
      ".\n"
     ]
    }
   ],
   "source": [
    "for tool_name in tools_path:\n",
    "    print(f\"Executing tool: {tool_name}\")\n",
    "    current_output = await execute_tool(\n",
    "        tool_name, \n",
    "        question, \n",
    "        question, \n",
    "        file_path_list, \n",
    "        config.MODEL_NAME, \n",
    "        kg_table_retriever, \n",
    "        total_output\n",
    "        )\n",
    "    print(f\"\\ncurrent_output: {current_output}\\n\")\n",
    "    total_output += current_output\n",
    "    responses_content += current_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    }
   ],
   "source": [
    "answer_text = await generate_output(question, total_output, model=config.MODEL_NAME)\n",
    "is_completed, current_final_answer = extract_answer_parts(answer_text)\n",
    "responses_content += f\"\\nTask completion status: {is_completed}\\nFinal Answer: {current_final_answer}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed\n",
      "Yes, the problem has been solved. The performance of three different machine learning algorithms (MLP, AdaBoost, and Random Forest) for predicting reactivity using electrical descriptors has been compared. The test accuracies are as follows:\n",
      "\n",
      "- MLP Classifier: Test Accuracy = 0.9545\n",
      "- AdaBoost Classifier: Test Accuracy = 0.8636\n",
      "- Random Forest Classifier: Test Accuracy = 0.9545\n",
      "\n",
      "Both MLP and Random Forest classifiers yield the best test accuracy of 0.9545.\n"
     ]
    }
   ],
   "source": [
    "print(is_completed)\n",
    "print(current_final_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Case III: Chemical synthesis and analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PROJECT_ROOT: /home/hjj/project/SciToolAgent\n",
      "PERSIST_DIR: /home/hjj/project/SciToolAgent/KG/storage_graph_large\n",
      "DATA_FILE_PATH: /home/hjj/project/SciToolAgent/data/toolsKG_full.xlsx\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "PROJECT_ROOT = os.getcwd()\n",
    "\n",
    "load_dotenv(os.path.join(PROJECT_ROOT, 'example.env'), override=True)\n",
    "\n",
    "PERSIST_DIR = os.path.join(PROJECT_ROOT, os.getenv(\"PERSIST_DIR\"))\n",
    "DATA_FILE_PATH = os.path.join(PROJECT_ROOT, os.getenv(\"DATA_FILE_PATH\"))\n",
    "\n",
    "print(\"PROJECT_ROOT:\", PROJECT_ROOT)\n",
    "print(\"PERSIST_DIR:\", PERSIST_DIR)\n",
    "print(\"DATA_FILE_PATH:\", DATA_FILE_PATH)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hjj/miniconda3/envs/toolsKG/lib/python3.9/site-packages/Bio/pairwise2.py:278: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "import logging\n",
    "import os\n",
    "import pdb\n",
    "import sys\n",
    "\n",
    "from scripts.custom_kg_retrievers import KGTableRetriever\n",
    "from scripts.retrieve_tool_info import execute_tool, check_toxicity\n",
    "from scripts.utils import load_index_from_storage_dir, construct_question\n",
    "from app.llms.planning import generate_tools_plan, generate_plan_input, generate_output\n",
    "from app.llms.extraction import extract_tool_path, extract_answer_parts\n",
    "from app.tools.safety.toxicity_checker import check_toxicity_by_type\n",
    "from app.tools.safety import extract_smiles_or_sequences\n",
    "from app.tools import run_tool\n",
    "from app.core.config import Config\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Settings\n",
    "from llama_index.core import QueryBundle\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "# print()\n",
    "config = Config()\n",
    "Settings.embed_model = OpenAIEmbedding(model=config.EMBEDDING_MODEL_NAME)\n",
    "Settings.llm = OpenAI(temperature=0, model=config.MODEL_NAME)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "question = \"Please predict the product of the chemical reaction between OC(=O)CCC(=O)O and O=C(C)OC(=O)C, and provide a textual description of the corresponding SELFIES. Moreover, check whether the substance has been patented and if it is explosive through CAS.\"\n",
    "is_add_filename = False\n",
    "file_path_list = []\n",
    "\n",
    "question = construct_question(question, file_path_list, is_add_filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.core.indices.loading:Loading all indices.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Knowledge Graph Index loaded successfully from KG/storage_graph_large\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "['CheckExplosiveness', 'TexToMoleculeSELFIES', 'RXNRetrosynthetic', 'GenerateMoleculeDescription', 'RXNPredict', 'SMILEStoSELFIES', 'Convert3DMolecules2SMILES', 'InChIKeyToSMILES', 'SELFIEStoSMILES', 'CheckPatent']\n"
     ]
    }
   ],
   "source": [
    "responses_content = \"\"\n",
    "total_output = \"\"\n",
    "\n",
    "storage_context, kg_index = load_index_from_storage_dir(config.PERSIST_DIR)\n",
    "kg_table_retriever = KGTableRetriever(\n",
    "    index=kg_index, \n",
    "    similarity_top_k=5, \n",
    "    graph_store_query_depth=3, \n",
    "    max_tools=10\n",
    ")        \n",
    "query_bundle = QueryBundle(query_str=question)\n",
    "tools_retrieved = await kg_table_retriever._retrieve(query_bundle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Request URL: https://api.shubiaobiao.com/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "['RXNPredict', 'SMILEStoSELFIES', 'GenerateMoleculeDescription', 'CheckPatent', 'CheckExplosiveness']\n"
     ]
    }
   ],
   "source": [
    "tools_plan = await generate_tools_plan(\n",
    "    question=question, \n",
    "    tools_retrieved_list=tools_retrieved, \n",
    "    model=config.MODEL_NAME, \n",
    "    previous_answer=total_output, \n",
    ")\n",
    "tools_path = extract_tool_path(tools_plan)\n",
    "print(tools_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing tool: RXNPredict\n",
      "Request URL: https://api.shubiaobiao.com/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: OC(=O)CCC(=O)O.O=C(C)OC(=O)C\n",
      "current_output: Now we have run the RXNPredict and obtained result \n",
      "### Reaction Prediction\n",
      "\n",
      "#### Reactants\n",
      "\n",
      "- **Reactants**: `OC(=O)CCC(=O)O.O=C(C)OC(=O)C`\n",
      "\n",
      "#### Predicted Product\n",
      "\n",
      "- **Product**: `O=C1CCC(=O)O1`\n",
      "\n",
      "#### More Information\n",
      "\n",
      "- [IBM RXN for Chemistry](https://rxn.res.ibm.com)\n",
      ".\n",
      "Executing tool: SMILEStoSELFIES\n",
      "Request URL: https://api.shubiaobiao.com/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: O=C1CCC(=O)O1\n",
      "current_output: Now we have run the SMILEStoSELFIES and obtained result \n",
      "**SMILES to SELFIES**\n",
      "Input: O=C1CCC(=O)O1\n",
      "\n",
      "**Reuslt**\n",
      "[O][=C][C][C][C][=Branch1][C][=O][O][Ring1][=Branch1]\n",
      ".\n",
      "Executing tool: GenerateMoleculeDescription\n",
      "Request URL: https://api.shubiaobiao.com/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: [O][=C][C][C][C][=Branch1][C][=O][O][Ring1][=Branch1]\n",
      "current_output: Now we have run the GenerateMoleculeDescription and obtained result **SELFIES:** `[O][=C][C][C][C][=Branch1][C][=O][O][Ring1][=Branch1]`\n",
      "\n",
      "**Description:** The molecule is a tetrahydrofurandione that is succinic anhydride substituted by oxo groups at positions 2 and 5. It has a role as a metabolite..\n",
      "Executing tool: CheckPatent\n",
      "Request URL: https://api.shubiaobiao.com/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: O=C1CCC(=O)O1\n",
      "current_output: Now we have run the CheckPatent and obtained result \n",
      "### Molecule: O=C1CCC(=O)O1\n",
      "\n",
      "#### Patent Status\n",
      "\n",
      "- **Status**: `Novel`\n",
      "\n",
      "#### More Information\n",
      "\n",
      "- [SureChEMBL](https://www.surechembl.org/search/?query=O=C1CCC(=O)O1)\n",
      ".\n",
      "Executing tool: CheckExplosiveness\n",
      "Request URL: https://api.shubiaobiao.com/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: O=C1CCC(=O)O1\n",
      "current_output: Now we have run the CheckExplosiveness and obtained result \n",
      "### Explosiveness Check\n",
      "\n",
      "#### CAS Number\n",
      "\n",
      "- **CAS**: `O=C1CCC(=O)O1`\n",
      "\n",
      "#### Result\n",
      "\n",
      "- **Explosiveness**: The molecule is not explosive\n",
      ".\n"
     ]
    }
   ],
   "source": [
    "for tool_name in tools_path:\n",
    "    print(f\"Executing tool: {tool_name}\")\n",
    "    current_output = await execute_tool(\n",
    "        tool_name, \n",
    "        question, \n",
    "        question, \n",
    "        file_path_list, \n",
    "        config.MODEL_NAME, \n",
    "        kg_table_retriever, \n",
    "        total_output\n",
    "        )\n",
    "    print(f\"current_output: {current_output}\")\n",
    "    total_output += current_output\n",
    "    responses_content += current_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Request URL: https://api.shubiaobiao.com/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.shubiaobiao.com/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Completed'"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "answer_text = await generate_output(question, total_output, model=config.MODEL_NAME)\n",
    "is_completed, current_final_answer = extract_answer_parts(answer_text)\n",
    "responses_content += f\"\\nTask completion status: {is_completed}\\nFinal Answer: {current_final_answer}\"\n",
    "is_completed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your problem has been solved. The predicted product of the chemical reaction between OC(=O)CCC(=O)O and O=C(C)OC(=O)C is `O=C1CCC(=O)O1`. The corresponding SELFIES is `[O][=C][C][C][C][=Branch1][C][=O][O][Ring1][=Branch1]`, which describes the molecule as a tetrahydrofurandione that is succinic anhydride substituted by oxo groups at positions 2 and 5, functioning as a metabolite. The molecule has been checked for patents and is considered novel according to SureChEMBL. Additionally, the CAS number indicates that the molecule is not explosive.\n"
     ]
    }
   ],
   "source": [
    "print(current_final_answer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Case IV: MOF materials analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "\n",
    "PROJECT_ROOT = os.getcwd()\n",
    "load_dotenv(os.path.join(PROJECT_ROOT, 'example.env'), override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hjj/miniconda3/envs/toolsKG/lib/python3.9/site-packages/Bio/pairwise2.py:278: BiopythonDeprecationWarning: Bio.pairwise2 has been deprecated, and we intend to remove it in a future release of Biopython. As an alternative, please consider using Bio.Align.PairwiseAligner as a replacement, and contact the Biopython developers if you still need the Bio.pairwise2 module.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from scripts.custom_kg_retrievers import KGTableRetriever\n",
    "from scripts.retrieve_tool_info import execute_tool, check_toxicity\n",
    "from scripts.utils import load_index_from_storage_dir, construct_question\n",
    "from app.llms.planning import generate_tools_plan, generate_plan_input, generate_output\n",
    "from app.llms.extraction import extract_tool_path, extract_answer_parts\n",
    "from app.tools.safety.toxicity_checker import check_toxicity_by_type\n",
    "from app.tools.safety import extract_smiles_or_sequences\n",
    "from app.tools import run_tool\n",
    "from app.core.config import Config\n",
    "\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.core import Settings\n",
    "from llama_index.core import QueryBundle\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "# print()\n",
    "config = Config()\n",
    "Settings.embed_model = OpenAIEmbedding(model=config.EMBEDDING_MODEL_NAME)\n",
    "Settings.llm = OpenAI(temperature=0, model=config.MODEL_NAME)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I want to know the thermal stability and CO2 absorption of this mof material, and then convert the smiles of this mof material into cas to query the corresponding price Here is the list of files: HKUST-1.cif\n"
     ]
    }
   ],
   "source": [
    "question = \"I want to know the thermal stability and CO2 absorption of this mof material, and then convert the smiles of this mof material into cas to query the corresponding price.\"\n",
    "\n",
    "is_add_filename = True\n",
    "file_path_list = [\n",
    "            \"/home/hjj/project/ToolsKG/TempFiles/cif/HKUST-1.cif\", \n",
    "        ]\n",
    "\n",
    "question = construct_question(question, file_path_list, is_add_filename)\n",
    "print(question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:llama_index.core.indices.loading:Loading all indices.\n",
      "Knowledge Graph Index loaded successfully from KG/storage_graph_large\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/embeddings \"HTTP/1.1 200 OK\"\n",
      "['PredictStability', 'PredictAdsorption', 'GetFloatingSolventMolecules', 'MofLattice', 'MofFractionalCoordinates', 'GetTerminalIndices', 'MOFToSMILES', 'CASToPrice', 'SMILESToCAS', 'MofNeighborIndices']\n"
     ]
    }
   ],
   "source": [
    "\n",
    "storage_context, kg_index = load_index_from_storage_dir(config.PERSIST_DIR)\n",
    "kg_table_retriever = KGTableRetriever(\n",
    "    index=kg_index, \n",
    "    similarity_top_k=10, \n",
    "    graph_store_query_depth=5, \n",
    "    max_tools=10\n",
    ")        \n",
    "query_bundle = QueryBundle(query_str=question)\n",
    "tools_retrieved = await kg_table_retriever._retrieve(query_bundle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "https://api.ai-gaochao.cn/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "['PredictStability', 'PredictAdsorption', 'MOFToSMILES', 'SMILESToCAS', 'CASToPrice']\n"
     ]
    }
   ],
   "source": [
    "responses_content = \"\"\n",
    "total_output = \"\"\n",
    "\n",
    "tools_plan = await generate_tools_plan(\n",
    "    question=question, \n",
    "    tools_retrieved_list=tools_retrieved, \n",
    "    model=config.MODEL_NAME, \n",
    "    previous_answer=total_output, \n",
    ")\n",
    "tools_path = extract_tool_path(tools_plan)\n",
    "print(tools_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing tool: PredictStability\n",
      "https://api.ai-gaochao.cn/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: HKUST-1.cif\n",
      "current_output: \n",
      "\n",
      "Now we have run the **PredictStability** and obtained result\n",
      "    filename  ANN_thermal_prediction  ANN_solvent_removal_prediction\n",
      "HKUST-1.cif              324.049896                        0.948832.\n",
      "Executing tool: PredictAdsorption\n",
      "https://api.ai-gaochao.cn/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: HKUST-1.cif\n",
      "current_output: \n",
      "\n",
      "Now we have run the **PredictAdsorption** and obtained result\n",
      "    filename  finished  CO2_absolute_mg/g  H2_absolute_mg/g  CO2_excess_mg/g  H2_excess_mg/g\n",
      "HKUST-1.cif      True         106.885632          0.052386       106.885632        0.052386.\n",
      "Executing tool: MOFToSMILES\n",
      "https://api.ai-gaochao.cn/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: HKUST-1.cif\n",
      "current_output: \n",
      "\n",
      "Now we have run the **MOFToSMILES** and obtained result\n",
      " | Name | SMILES | Status |\n",
      "|------|--------|--------|\n",
      "| HKUST-1 | [O]C(=O)c1cc(C([O])=O)cc(C([O])=O)c1 | Success |\n",
      ".\n",
      "Executing tool: SMILESToCAS\n",
      "https://api.ai-gaochao.cn/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: [O]C(=O)c1cc(C([O])=O)cc(C([O])=O)c1\n",
      "current_output: \n",
      "\n",
      "Now we have run the **SMILESToCAS** and obtained result\n",
      " #### Smiles \n",
      " [O]C(=O)c1cc(C([O])=O)cc(C([O])=O)c1\n",
      " #### CAS \n",
      " 554-95-0.\n",
      "Executing tool: CASToPrice\n",
      "https://api.ai-gaochao.cn/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n",
      "input_str: 554-95-0\n",
      "current_output: \n",
      "\n",
      "Now we have run the **CASToPrice** and obtained result\n",
      " #### CAS \n",
      "554-95-0\n",
      " #### Average price \n",
      " ￥0.4096666666666667.\n"
     ]
    }
   ],
   "source": [
    "for tool_name in tools_path:\n",
    "    print(f\"Executing tool: {tool_name}\")\n",
    "    current_output = await execute_tool(\n",
    "        tool_name, \n",
    "        question, \n",
    "        question, \n",
    "        file_path_list, \n",
    "        config.MODEL_NAME, \n",
    "        kg_table_retriever, \n",
    "        total_output\n",
    "    )\n",
    "    print(f\"current_output: {current_output}\")\n",
    "    total_output += current_output\n",
    "    responses_content += current_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "https://api.ai-gaochao.cn/v1/chat/completions\n",
      "INFO:httpx:HTTP Request: POST https://api.ai-gaochao.cn/v1/chat/completions \"HTTP/1.1 200 OK\"\n"
     ]
    }
   ],
   "source": [
    "answer_text = await generate_output(question, total_output, model=config.MODEL_NAME)\n",
    "is_completed, current_final_answer = extract_answer_parts(answer_text)\n",
    "responses_content += f\"\\nTask completion status: {is_completed}\\nFinal Answer: {current_final_answer}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed\n",
      "1. Thermal Stability of HKUST-1: 324.049896 Kelvin\n",
      "2. CO2 Absorption of HKUST-1: 106.885632 mg/g\n",
      "3. SMILES to CAS Conversion: The SMILES [O]C(=O)c1cc(C([O])=O)cc(C([O])=O)c1 corresponds to CAS number 554-95-0.\n",
      "4. Price Query: The average price for CAS number 554-95-0 is ￥0.4096666666666667.\n"
     ]
    }
   ],
   "source": [
    "print(is_completed)\n",
    "print(current_final_answer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "toolsKG",
   "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}