{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# o1 Style Thinking\n",
    "\n",
    "In this recipe, we will show how to achieve Chain-of-Thought Reasoning.\n",
    "This makes LLMs to breakdown the task in multiple steps and generate a coherent output allowing to solve complex tasks in logical steps.\n",
    "\n",
    "<div class=\"admonition tip\">\n",
    "<p class=\"admonition-title\">Mirascope Concepts Used</p>\n",
    "<ul>\n",
    "<li><a href=\"../../../learn/prompts/\">Prompts</a></li>\n",
    "<li><a href=\"../../../learn/calls/\">Calls</a></li>\n",
    "<li><a href=\"../../../learn/response_models/\">Response Models</a></li>\n",
    "</ul>\n",
    "</div>\n",
    "\n",
    "<div class=\"admonition note\">\n",
    "<p class=\"admonition-title\">Background</p>\n",
    "<p>\n",
    "    Large Language Models (LLMs) are known to generate text that is coherent and fluent. However, they often struggle with tasks that require multi-step reasoning or logical thinking. In this recipe, we will show how to use Mirascope to guide the LLM to break down the task into multiple steps and generate a coherent output.\n",
    "\n",
    "</p>\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "To set up our environment, first let's install all of the packages we will use:\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"mirascope[groq]\" \n",
    "!pip install datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "# Set the appropriate API key for the provider you're using\n",
    "# Here we are using GROQ_API_KEY\n",
    "\n",
    "export GROQ_API_KEY=\"Your API Key\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Without Chain-of-Thought Reasoning\n",
    "\n",
    "We will begin by showing how a typical LLM performs on a task that requires multi-step reasoning. In this example, we will ask the model to generate a count the number of `s`s in the word `Mississssippi` (Yes it has 7`s`'s). We will use the `llama-3.1-8b-instant` for this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(User): how many s's in the word mississssippi\n",
      "(Assistant): There are 5 s's in the word \"Mississippi\".\n"
     ]
    }
   ],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "from mirascope.core import groq\n",
    "\n",
    "history: list[dict[str, str]] = []\n",
    "\n",
    "\n",
    "@groq.call(\"llama-3.1-8b-instant\")\n",
    "def generate_answer(question: str) -> str:\n",
    "    return f\"Generate an answer to this question: {question}\"\n",
    "\n",
    "\n",
    "def run() -> None:\n",
    "    question: str = \"how many s's in the word mississssippi\"\n",
    "    response: str = generate_answer(question)\n",
    "    print(f\"(User): {question}\")\n",
    "    print(f\"(Assistant): {response}\")\n",
    "    history.append({\"role\": \"user\", \"content\": question})\n",
    "    history.append({\"role\": \"assistant\", \"content\": response})\n",
    "\n",
    "\n",
    "run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, the zero-shot method is used to generate the output. The model is not provided with any additional information or context to help it generate the output. The model is only given the input prompt and asked to generate the output.\n",
    "\n",
    "This is not so effective when there is a logcial task to be performed.\n",
    "\n",
    "Now let's see how the model performs on this task when it can reason using Chain-of-Thought Reasoning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# With Chain of Thought Reasoning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(User): How many s's are in the word 'mississssippi'?\n",
      "Step 1: Initial Assessment and Counting:\n",
      "To count the number of 's's in the word 'mississssippi', I will first notice that the 's's appear together in two groups. This makes it easier to count. The first group contains 1 's', and the second group contains 4 's's. Additionally, there is 1 more 's' separate from these groups. By combining the counts from these groups and the additional 's', I arrive at a preliminary total of 6 's's. However, upon reviewing the options, I realize that I must consider the possibility of an error in my initial assessment.\n",
      "**Thinking time: 1.13 seconds**\n",
      "\n",
      "Step 2: Re-examining the Count of 's's in the Word 'mississssippi':\n",
      "I will recount the 's's in the word 'mississssippi'. Upon re-examination, I notice the groups of 's's are actually 'm-i-s-s-i-s-sss-ss-i-pp-i'. Recounting the groups, I still find two 's's separate from the groups and 4 in the last group. I still find 1 's' in a separate group at the start. Combining them I get 7 's's in the word. This seems correct, however, I must explore any possible alternative count. In considering my count, I consider the alternative the individual 's's in 'mississssippi' are not as grouped, but separate. I manually recount: the first 's', then 4 's's, plus 2 's's at the end of the groups gives me 7. Upon consideration, this approach still indicates there are 7 's's.\n",
      "**Thinking time: 1.21 seconds**\n",
      "\n",
      "Step 3: Validating the Count of 's's in the Word 'mississssippi':\n",
      "Considering the count I arrived at in the previous steps, I notice that I must further ensure the count is correct. There are no other apparent alternative methods to consider. Upon reflection on my approach, my confidence in my methods is high enough to proceed. However, this confidence does not exclude the possibility of an error.\n",
      "**Thinking time: 0.74 seconds**\n",
      "\n",
      "Final Answer:\n",
      "There are 7 's's in the word 'mississssippi'.\n",
      "**Thinking time: 0.42 seconds**\n",
      "\n",
      "**Total thinking time: 3.50 seconds**\n"
     ]
    }
   ],
   "source": [
    "from typing import Literal\n",
    "\n",
    "from mirascope.core import groq\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "history: list[dict] = []\n",
    "\n",
    "\n",
    "class COTResult(BaseModel):\n",
    "    title: str = Field(..., desecription=\"The title of the step\")\n",
    "    content: str = Field(..., description=\"The output content of the step\")\n",
    "    next_action: Literal[\"continue\", \"final_answer\"] = Field(\n",
    "        ..., description=\"The next action to take\"\n",
    "    )\n",
    "\n",
    "\n",
    "@groq.call(\"llama-3.1-70b-versatile\", json_mode=True, response_model=COTResult)\n",
    "def cot_step(prompt: str, step_number: int, previous_steps: str) -> str:\n",
    "    return f\"\"\"\n",
    "    You are an expert AI assistant that explains your reasoning step by step.\n",
    "    For this step, provide a title that describes what you're doing, along with the content.\n",
    "    Decide if you need another step or if you're ready to give the final answer.\n",
    "\n",
    "    Guidelines:\n",
    "    - Use AT MOST 5 steps to derive the answer.\n",
    "    - Be aware of your limitations as an LLM and what you can and cannot do.\n",
    "    - In your reasoning, include exploration of alternative answers.\n",
    "    - Consider you may be wrong, and if you are wrong in your reasoning, where it would be.\n",
    "    - Fully test all other possibilities.\n",
    "    - YOU ARE ALLOWED TO BE WRONG. When you say you are re-examining\n",
    "        - Actually re-examine, and use another approach to do so.\n",
    "        - Do not just say you are re-examining.\n",
    "\n",
    "    IMPORTANT: Do not use code blocks or programming examples in your reasoning. Explain your process in plain language.\n",
    "\n",
    "    This is step number {step_number}.\n",
    "\n",
    "    Question: {prompt}\n",
    "\n",
    "    Previous steps:\n",
    "    {previous_steps}\n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "@groq.call(\"llama-3.1-70b-versatile\")\n",
    "def final_answer(prompt: str, reasoning: str) -> str:\n",
    "    return f\"\"\"\n",
    "    Based on the following chain of reasoning, provide a final answer to the question.\n",
    "    Only provide the text response without any titles or preambles.\n",
    "    Retain any formatting as instructed by the original prompt, such as exact formatting for free response or multiple choice.\n",
    "\n",
    "    Question: {prompt}\n",
    "\n",
    "    Reasoning:\n",
    "    {reasoning}\n",
    "\n",
    "    Final Answer:\n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "def generate_cot_response(\n",
    "    user_query: str,\n",
    ") -> tuple[list[tuple[str, str, float]], float]:\n",
    "    steps: list[tuple[str, str, float]] = []\n",
    "    total_thinking_time: float = 0.0\n",
    "    step_count: int = 1\n",
    "    reasoning: str = \"\"\n",
    "    previous_steps: str = \"\"\n",
    "\n",
    "    while True:\n",
    "        start_time: datetime = datetime.now()\n",
    "        cot_result = cot_step(user_query, step_count, previous_steps)\n",
    "        end_time: datetime = datetime.now()\n",
    "        thinking_time: float = (end_time - start_time).total_seconds()\n",
    "\n",
    "        steps.append(\n",
    "            (\n",
    "                f\"Step {step_count}: {cot_result.title}\",\n",
    "                cot_result.content,\n",
    "                thinking_time,\n",
    "            )\n",
    "        )\n",
    "        total_thinking_time += thinking_time\n",
    "\n",
    "        reasoning += f\"\\n{cot_result.content}\\n\"\n",
    "        previous_steps += f\"\\n{cot_result.content}\\n\"\n",
    "\n",
    "        if cot_result.next_action == \"final_answer\" or step_count >= 5:\n",
    "            break\n",
    "\n",
    "        step_count += 1\n",
    "\n",
    "    # Generate final answer\n",
    "    start_time = datetime.now()\n",
    "    final_result: str = final_answer(user_query, reasoning).content\n",
    "    end_time = datetime.now()\n",
    "    thinking_time = (end_time - start_time).total_seconds()\n",
    "    total_thinking_time += thinking_time\n",
    "\n",
    "    steps.append((\"Final Answer\", final_result, thinking_time))\n",
    "\n",
    "    return steps, total_thinking_time\n",
    "\n",
    "\n",
    "def display_cot_response(\n",
    "    steps: list[tuple[str, str, float]], total_thinking_time: float\n",
    ") -> None:\n",
    "    for title, content, thinking_time in steps:\n",
    "        print(f\"{title}:\")\n",
    "        print(content.strip())\n",
    "        print(f\"**Thinking time: {thinking_time:.2f} seconds**\\n\")\n",
    "\n",
    "    print(f\"**Total thinking time: {total_thinking_time:.2f} seconds**\")\n",
    "\n",
    "\n",
    "def run() -> None:\n",
    "    question: str = \"How many s's are in the word 'mississssippi'?\"\n",
    "    print(\"(User):\", question)\n",
    "    # Generate COT response\n",
    "    steps, total_thinking_time = generate_cot_response(question)\n",
    "    display_cot_response(steps, total_thinking_time)\n",
    "\n",
    "    # Add the interaction to the history\n",
    "    history.append({\"role\": \"user\", \"content\": question})\n",
    "    history.append(\n",
    "        {\"role\": \"assistant\", \"content\": steps[-1][1]}\n",
    "    )  # Add only the final answer to the history\n",
    "\n",
    "\n",
    "# Run the function\n",
    "\n",
    "run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As demonstrated in the COT Reasoning example, we can guide the model to break down the task into multiple steps and generate a coherent output. This allows the model to solve complex tasks in logical steps.\n",
    "However, this requires multiple calls to the model, which may be expensive in terms of cost and time.\n",
    "Also model may not always identify the correct steps to solve the task, hence is not deterministic."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "Chain of Thought Reasoning is a powerful technique that allows LLMs to solve complex tasks in logical steps. However, it requires multiple calls to the model and may not always identify the correct steps to solve the task. This technique can be useful when the task requires multi-step reasoning or logical thinking.\n",
    "\n",
    "Care should be taken to ensure that the model is guided correctly and that the output is coherent and accurate."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}