{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-Agent Question-Answer\n",
    "\n",
    "In this example, we are building a workflow for [HotpotQA](https://arxiv.org/abs/1809.09600), which requires the agent to retrieve from wiki-2017 documents twice to answer a factorial question.\n",
    "\n",
    "The implementation is adopted from [dspy](https://github.com/stanfordnlp/dspy?tab=readme-ov-file#5a-dspy-vs-thin-wrappers-for-prompts-openai-api-minichain-basic-templating), including three agents in total:\n",
    "- **Query agent 0**: generates a search query from the user question.\n",
    "- **Query agent 1**: refines the search by retrieving additional information based on initial results.\n",
    "- **Answer agent**: synthesizes the retrieved documents to provide a final answer.\n",
    "\n",
    "![hotpotqa](../imgs/hotpotqa.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) Setup\n",
    "\n",
    "First, let's set the environment for workflow execution. Following keys are required:\n",
    "\n",
    "OPENAI_API_KEY=\"your-openai-key\"\n",
    "COLBERT_URL=\"colbert-serving-url\"\n",
    "\n",
    "> **Note:** \n",
    ">\n",
    "> If you are using DSPy's ColBERT service, try link `http://20.102.90.50:2017/wiki17_abstracts`. \n",
    ">\n",
    "> For hosting on your local machine, check [ColBERT official repo](https://github.com/stanford-futuredata/ColBERT) for installation and setup."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2) Check HotPotQA Workflow\n",
    "\n",
    "The complete code for this workflow is based on `dspy` and is avaibale in `workflow.py`. Try it out with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'answer': 'The 2010 population of Woodmere, New York, the birthplace of Gerard Piel, was 17,121.'}\n"
     ]
    }
   ],
   "source": [
    "%run workflow.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3) Optimize The Workflow\n",
    "\n",
    "The workflow entry point is already registered using annotation `cognify.register_workflow`.\n",
    "\n",
    "Here we configure the optimization pipeline:\n",
    "1. Define the evaluation method\n",
    "2. Define the data loader\n",
    "3. Config the optimizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Tell Cognify how good the answer is\n",
    "\n",
    "We use builtin f1 score to evaluate the similarity between the predicted answer and the given ground truth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cognify\n",
    "from cognify.hub.evaluators import f1_score_str\n",
    "\n",
    "@cognify.register_evaluator\n",
    "def answer_f1(answer: str, ground_truth: str):\n",
    "    return f1_score_str(answer, ground_truth)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Tell Cognify what data to use\n",
    "\n",
    "We directly use the hotpotqa dataset from DSPy with some minor formatting changes.\n",
    "\n",
    "The loaded data should be a series of pairs (input / ground_truth). \n",
    "\n",
    "Both `input` and `ground_truth` should be a dictionary.\n",
    "\n",
    "Cognify will dispath the data by matching their name to the function signature, in short:\n",
    "\n",
    "```python\n",
    "# register workflow\n",
    "# register evaluator\n",
    "\n",
    "data: [(input, ground_truth), ...] = data_loader()\n",
    "for input, ground_truth in data:\n",
    "    result = workflow(**input)\n",
    "    eval_inputs = as_per_func_signature(evaluator, input, result, ground_truth)\n",
    "    score = evaluator(**eval_inputs)\n",
    "```\n",
    "\n",
    "According to the above rule, we register the data loader as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def formatting(item):\n",
    "    return (\n",
    "        {'question': item.question},\n",
    "        {'ground_truth': item.answer}\n",
    "    )\n",
    "\n",
    "@cognify.register_data_loader\n",
    "def load_hotpotqa_data():\n",
    "    from dspy.datasets.hotpotqa import HotPotQA\n",
    "    dataset = HotPotQA(train_seed=1, train_size=150, eval_seed=2023, dev_size=200, test_size=0)\n",
    "    \n",
    "    trainset = [formatting(x) for x in dataset.train[0:100]]\n",
    "    valset = [formatting(x) for x in dataset.train[100:150]]\n",
    "    devset = [formatting(x) for x in dataset.dev]\n",
    "    return trainset, valset, devset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Config the optimizer\n",
    "\n",
    "Let's use the default configuration to optimize this workflow. The search space includes:\n",
    "- 2 fewshot examples to add for each agent\n",
    "- whether to apply Chain-of-thought to each agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cognify.hub.search import default\n",
    "\n",
    "search_settings = default.create_search()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Start the Optimization\n",
    "\n",
    "You can save the above configs in `config.py` file and use Cognify's CLI to fire the optimization with:\n",
    "\n",
    "```console\n",
    "$ cognify optimize workflow.py\n",
    "```\n",
    "\n",
    "Alternatively you can run the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, val, dev = load_hotpotqa_data()\n",
    "\n",
    "opt_cost, pareto_frontier, opt_logs = cognify.optimize(\n",
    "    script_path=\"workflow.py\",\n",
    "    control_param=search_settings,\n",
    "    train_set=train,\n",
    "    val_set=val,\n",
    "    eval_fn=answer_f1,\n",
    "    force=True, # This will overwrite the existing results\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Optimization Results\n",
    "\n",
    "Cognfiy will output each optimized workflow to a `.cog` file. For this workflow, the optimizer chooses the following optimizations:\n",
    "- ensemble the first query generation module\n",
    "- add few-shot examples to the ensembled query generation modules\n",
    "- for the answer generation module, add few-shot examples. \n",
    "\n",
    "The final optimized workflow is depicted below, with optimizations highlighted in green.\n",
    "\n",
    "![hotpotqa-opt](../imgs/hotpotqa_optimized.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The few-shot examples inserted into the prompt for the query generation modules were as follows:\n",
    "\n",
    "\n",
    "> **Demonstration 1**:  \n",
    "> Input (question): \"Gustav Mahler composed a beautiful piece performed by the Bach-Elgar Choir. What is the name of that piece??\"\n",
    "> Output (search query): \"Gustav Mahler piece performed by Bach-Elgar Choir\"\n",
    "> \n",
    "> **Demonstration 2**:  \n",
    "> Input (question): \"Merle Reagle did crosswords for what magazine that has a focus on aging issues?\"  \n",
    "> Output (search query): \"Merle Reagle crosswords magazine aging issues\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The few-shot examples inserted into the prompt for the answer generation modules were as follows:\n",
    "\n",
    "> **Demonstration 1:**  \n",
    "> \tInput (context): [\"Bach-Elgar Choir | The Bach-Elgar Choir is a community chorus of long standing in Hamilton, Ontario. The Choir is composed of accomplished amateur singers from Hamilton... Notable performances include ... and Mahler's 'Symphony No. 2'...\", \"Symphony No. 8 (Mahler) | The Symphony No. 8 in E-flat major by Gustav Mahler is one of the largest-scale choral works...\", ... ] (truncated for brevity)\n",
    "> \n",
    "> \tInput (question): \"Gustav Mahler composed a beautiful piece performed by the Bach-Elgar Choir. What is the name of that piece??\" \n",
    "> \t\n",
    "> \tOutput (answer): \"Symphony No. 2 (the Resurrection)\"\n",
    ">  \n",
    "> **Demonstration 2:**  \n",
    "> \tInput (context): [\"Merl Reagle | Merl Harry Reagle (January 5, 1950 \\u2013 August 22, 2015) was an American crossword constructor. For 30 years, he constructed a puzzle every Sunday...Reagle also produced a bimonthly crossword puzzle for 'AARP The Magazine' magazine...\", \"Aging and Disease | Aging and Disease is a bimonthly peer-reviewed open access medical journal...\", \"AARP The Magazine | AARP The Magazine is an American bi-monthly magazine, published by the American Association of Retired People, AARP, which focuses on aging issues.\"] (truncated for brevity)\n",
    "> \n",
    ">   Input (question): \"Merle Reagle did crosswords for what magazine that has a focus on aging issues?\"  \n",
    "> \t \n",
    "> \tOutput (answer): \"AARP The Magazine\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check out more details on [how to interpret optimization results](https://cognify-ai.readthedocs.io/en/latest/user_guide/tutorials/interpret.html#detailed-transformation-trace)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}