{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Notebook for running Chain-of-Thought with supporting context experiments "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys, os\n",
    "sys.path.append('..')\n",
    "root = '../root/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import joblib\n",
    "import numpy as np\n",
    "from agents import CoTAgent, ReflexionStrategy\n",
    "from util import summarize_trial, log_trial, save_agents"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load the HotPotQA Sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "hotpot = joblib.load('../data/hotpot-qa-distractor-sample.joblib').reset_index(drop = True)\n",
    "\n",
    "hotpot['supporting_paragraphs'] = None\n",
    "for ind, row in hotpot.iterrows():\n",
    "    supporting_articles = row['supporting_facts']['title']\n",
    "    articles = row['context']['title']\n",
    "    sentences = row['context']['sentences'] \n",
    "    supporting_paragraphs = []\n",
    "    for article in supporting_articles:\n",
    "        supporting_paragraph = ''.join(sentences[np.where(articles == article)][0])\n",
    "        supporting_paragraphs.append(supporting_paragraph)\n",
    "    supporting_paragraphs = '\\n\\n'.join(supporting_paragraphs)\n",
    "    hotpot.at[ind, 'supporting_paragraphs'] = supporting_paragraphs"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define the Reflexion Strategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "    NONE: No reflection\n",
      "    LAST_ATTEMPT: Use last reasoning trace in context \n",
      "    REFLEXION: Apply reflexion to the next reasoning trace \n",
      "    LAST_ATTEMPT_AND_REFLEXION: Use last reasoning trace in context and apply reflexion to the next reasoning trace \n",
      "    \n"
     ]
    }
   ],
   "source": [
    "print(ReflexionStrategy.__doc__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "strategy: ReflexionStrategy = ReflexionStrategy.REFLEXION"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize a CoTAgent for each question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prompts import cot_agent_prompt, cot_reflect_agent_prompt, cot_reflect_prompt\n",
    "from fewshots import COT, COT_REFLECT\n",
    "agents = [CoTAgent(row['question'],\n",
    "                   row['supporting_paragraphs'],\n",
    "                   row['answer'],\n",
    "                   agent_prompt=cot_agent_prompt if strategy == ReflexionStrategy.NONE else cot_reflect_agent_prompt,\n",
    "                   cot_examples=COT,\n",
    "                   reflect_prompt=cot_reflect_prompt,\n",
    "                   reflect_examples=COT_REFLECT,\n",
    "                    ) for _, row in hotpot.iterrows()]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run `n` trials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 5\n",
    "trial = 0\n",
    "log = ''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValidationError",
     "evalue": "1 validation error for HumanMessage\ncontent\n  field required (type=value_error.missing)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValidationError\u001b[0m                           Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(n):\n\u001b[1;32m      2\u001b[0m     \u001b[39mfor\u001b[39;00m agent \u001b[39min\u001b[39;00m [a \u001b[39mfor\u001b[39;00m a \u001b[39min\u001b[39;00m agents \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m a\u001b[39m.\u001b[39mis_correct()]:\n\u001b[0;32m----> 3\u001b[0m         agent\u001b[39m.\u001b[39;49mrun(reflexion_strategy \u001b[39m=\u001b[39;49m strategy)\n\u001b[1;32m      4\u001b[0m         \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39mAnswer: \u001b[39m\u001b[39m{\u001b[39;00magent\u001b[39m.\u001b[39mkey\u001b[39m}\u001b[39;00m\u001b[39m'\u001b[39m)\n\u001b[1;32m      5\u001b[0m     trial \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1\u001b[39m\n",
      "File \u001b[0;32m~/Documents/Research/reflexion/reflexion/hotpotqa_runs/notebooks/../agents.py:78\u001b[0m, in \u001b[0;36mCoTAgent.run\u001b[0;34m(self, reflexion_strategy)\u001b[0m\n\u001b[1;32m     76\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mreflect(reflexion_strategy)\n\u001b[1;32m     77\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mreset()\n\u001b[0;32m---> 78\u001b[0m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mstep()\n\u001b[1;32m     79\u001b[0m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mstep_n \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1\u001b[39m\n",
      "File \u001b[0;32m~/Documents/Research/reflexion/reflexion/hotpotqa_runs/notebooks/../agents.py:84\u001b[0m, in \u001b[0;36mCoTAgent.step\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m     81\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mstep\u001b[39m(\u001b[39mself\u001b[39m) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m     82\u001b[0m     \u001b[39m# Think\u001b[39;00m\n\u001b[1;32m     83\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mscratchpad \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39mf\u001b[39m\u001b[39m'\u001b[39m\u001b[39m\\n\u001b[39;00m\u001b[39mThought:\u001b[39m\u001b[39m'\u001b[39m\n\u001b[0;32m---> 84\u001b[0m     \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mscratchpad \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39m'\u001b[39m\u001b[39m \u001b[39m\u001b[39m'\u001b[39m \u001b[39m+\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mprompt_agent()\n\u001b[1;32m     85\u001b[0m     \u001b[39mprint\u001b[39m(\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mscratchpad\u001b[39m.\u001b[39msplit(\u001b[39m'\u001b[39m\u001b[39m\\n\u001b[39;00m\u001b[39m'\u001b[39m)[\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m])\n\u001b[1;32m     87\u001b[0m     \u001b[39m# Act\u001b[39;00m\n",
      "File \u001b[0;32m~/Documents/Research/reflexion/reflexion/hotpotqa_runs/notebooks/../agents.py:132\u001b[0m, in \u001b[0;36mCoTAgent.prompt_agent\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    131\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mprompt_agent\u001b[39m(\u001b[39mself\u001b[39m) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m \u001b[39mstr\u001b[39m:\n\u001b[0;32m--> 132\u001b[0m     \u001b[39mreturn\u001b[39;00m format_step(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49maction_llm(\u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_build_agent_prompt()))\n",
      "File \u001b[0;32m~/Documents/Research/reflexion/reflexion/hotpotqa_runs/notebooks/../llm.py:25\u001b[0m, in \u001b[0;36mAnyOpenAILLM.__call__\u001b[0;34m(self, prompt)\u001b[0m\n\u001b[1;32m     21\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel(prompt)\n\u001b[1;32m     22\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m     23\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel(\n\u001b[1;32m     24\u001b[0m         [\n\u001b[0;32m---> 25\u001b[0m             HumanMessage(\n\u001b[1;32m     26\u001b[0m                 context\u001b[39m=\u001b[39;49mprompt,\n\u001b[1;32m     27\u001b[0m             )\n\u001b[1;32m     28\u001b[0m         ]\n\u001b[1;32m     29\u001b[0m     )\u001b[39m.\u001b[39mcontent\n",
      "File \u001b[0;32m~/Documents/Research/reflexion/reflexion/env/lib/python3.11/site-packages/pydantic/main.py:341\u001b[0m, in \u001b[0;36mpydantic.main.BaseModel.__init__\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mValidationError\u001b[0m: 1 validation error for HumanMessage\ncontent\n  field required (type=value_error.missing)"
     ]
    }
   ],
   "source": [
    "for i in range(n):\n",
    "    for agent in [a for a in agents if not a.is_correct()]:\n",
    "        agent.run(reflexion_strategy = strategy)\n",
    "        print(f'Answer: {agent.key}')\n",
    "    trial += 1\n",
    "    log += log_trial(agents, trial)\n",
    "    correct, incorrect = summarize_trial(agents)\n",
    "    print(f'Finished Trial {trial}, Correct: {len(correct)}, Incorrect: {len(incorrect)}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save the result log"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(root, 'CoT', 'context', strategy.value, f'{len(agents)}_questions_{trial}_trials.txt'), 'w') as f:\n",
    "    f.write(log)\n",
    "save_agents(agents, os.path.join(root, 'CoT', 'context', strategy.value, 'agents'))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "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.11.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "e23f799cbd2581634725fbf6ce3480ae26192d78438dfafc8efe944acd6490d5"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}