{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a38e5d2d-7587-4192-90f2-b58e6c62f08c",
   "metadata": {},
   "source": [
    "# Self-Discover Agent\n",
    "\n",
    "An implementation of the [Self-Discover paper](https://arxiv.org/pdf/2402.03620.pdf).\n",
    "\n",
    "Based on [this implementation from @catid](https://github.com/catid/self-discover/tree/main?tab=readme-ov-file)\n",
    "\n",
    "\n",
    "## Define the prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a18d8f24-5d9a-45c5-9739-6f3c4ed6c9c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Self-Discovery Select Prompt:\n",
      "Select several reasoning modules that are crucial to utilize in order to solve the given task:\n",
      "\n",
      "All reasoning module descriptions:\n",
      "\u001b[33;1m\u001b[1;3m{reasoning_modules}\u001b[0m\n",
      "\n",
      "Task: \u001b[33;1m\u001b[1;3m{task_description}\u001b[0m\n",
      "\n",
      "Select several modules are crucial for solving the task above:\n",
      "\n",
      "Self-Discovery Select Response:\n",
      "Rephrase and specify each reasoning module so that it better helps solving the task:\n",
      "\n",
      "SELECTED module descriptions:\n",
      "\u001b[33;1m\u001b[1;3m{selected_modules}\u001b[0m\n",
      "\n",
      "Task: \u001b[33;1m\u001b[1;3m{task_description}\u001b[0m\n",
      "\n",
      "Adapt each reasoning module description to better solve the task:\n",
      "\n",
      "Self-Discovery Structured Prompt:\n",
      "Operationalize the reasoning modules into a step-by-step reasoning plan in JSON format:\n",
      "\n",
      "Here's an example:\n",
      "\n",
      "Example task:\n",
      "\n",
      "If you follow these instructions, do you return to the starting point? Always face forward. Take 1 step backward. Take 9 steps left. Take 2 steps backward. Take 6 steps forward. Take 4 steps forward. Take 4 steps backward. Take 3 steps right.\n",
      "\n",
      "Example reasoning structure:\n",
      "\n",
      "{\n",
      "    \"Position after instruction 1\":\n",
      "    \"Position after instruction 2\":\n",
      "    \"Position after instruction n\":\n",
      "    \"Is final position the same as starting position\":\n",
      "}\n",
      "\n",
      "Adapted module description:\n",
      "\u001b[33;1m\u001b[1;3m{adapted_modules}\u001b[0m\n",
      "\n",
      "Task: \u001b[33;1m\u001b[1;3m{task_description}\u001b[0m\n",
      "\n",
      "Implement a reasoning structure for solvers to follow step-by-step and arrive at correct answer.\n",
      "\n",
      "Note: do NOT actually arrive at a conclusion in this pass. Your job is to generate a PLAN so that in the future you can fill it out and arrive at the correct conclusion for tasks like this\n",
      "Self-Discovery Structured Response:\n",
      "Follow the step-by-step reasoning plan in JSON to correctly solve the task. Fill in the values following the keys by reasoning specifically about the task given. Do not simply rephrase the keys.\n",
      "    \n",
      "Reasoning Structure:\n",
      "\u001b[33;1m\u001b[1;3m{reasoning_structure}\u001b[0m\n",
      "\n",
      "Task: \u001b[33;1m\u001b[1;3m{task_description}\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from langchain import hub\n",
    "\n",
    "select_prompt = hub.pull(\"hwchase17/self-discovery-select\")\n",
    "print(\"Self-Discovery Select Prompt:\")\n",
    "select_prompt.pretty_print()\n",
    "print(\"Self-Discovery Select Response:\")\n",
    "adapt_prompt = hub.pull(\"hwchase17/self-discovery-adapt\")\n",
    "adapt_prompt.pretty_print()\n",
    "structured_prompt = hub.pull(\"hwchase17/self-discovery-structure\")\n",
    "print(\"Self-Discovery Structured Prompt:\")\n",
    "structured_prompt.pretty_print()\n",
    "reasoning_prompt = hub.pull(\"hwchase17/self-discovery-reasoning\")\n",
    "print(\"Self-Discovery Structured Response:\")\n",
    "reasoning_prompt.pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bce1135e",
   "metadata": {},
   "source": [
    "## Define the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9f554045-6e79-42d3-be4b-835bbbd0b78c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional, TypedDict\n",
    "\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "from langgraph.graph import END, START, StateGraph\n",
    "\n",
    "\n",
    "class SelfDiscoverState(TypedDict):\n",
    "    reasoning_modules: str\n",
    "    task_description: str\n",
    "    selected_modules: Optional[str]\n",
    "    adapted_modules: Optional[str]\n",
    "    reasoning_structure: Optional[str]\n",
    "    answer: Optional[str]\n",
    "\n",
    "\n",
    "model = ChatOpenAI(temperature=0, model=\"gpt-4-turbo-preview\")\n",
    "\n",
    "\n",
    "def select(inputs):\n",
    "    select_chain = select_prompt | model | StrOutputParser()\n",
    "    return {\"selected_modules\": select_chain.invoke(inputs)}\n",
    "\n",
    "\n",
    "def adapt(inputs):\n",
    "    adapt_chain = adapt_prompt | model | StrOutputParser()\n",
    "    return {\"adapted_modules\": adapt_chain.invoke(inputs)}\n",
    "\n",
    "\n",
    "def structure(inputs):\n",
    "    structure_chain = structured_prompt | model | StrOutputParser()\n",
    "    return {\"reasoning_structure\": structure_chain.invoke(inputs)}\n",
    "\n",
    "\n",
    "def reason(inputs):\n",
    "    reasoning_chain = reasoning_prompt | model | StrOutputParser()\n",
    "    return {\"answer\": reasoning_chain.invoke(inputs)}\n",
    "\n",
    "\n",
    "graph = StateGraph(SelfDiscoverState)\n",
    "graph.add_node(select)\n",
    "graph.add_node(adapt)\n",
    "graph.add_node(structure)\n",
    "graph.add_node(reason)\n",
    "graph.add_edge(START, \"select\")\n",
    "graph.add_edge(\"select\", \"adapt\")\n",
    "graph.add_edge(\"adapt\", \"structure\")\n",
    "graph.add_edge(\"structure\", \"reason\")\n",
    "graph.add_edge(\"reason\", END)\n",
    "app = graph.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29fe385b-cf5d-4581-80e7-55462f5628bb",
   "metadata": {},
   "source": [
    "## Invoke the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6cbfbe81-f751-42da-843a-f9003ace663d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'select': {'selected_modules': 'To solve the task of identifying the shape drawn by the SVG path element, the following reasoning modules are crucial:\\n\\n1. **Critical Thinking (10):** This involves analyzing the provided SVG path commands to understand how they contribute to forming a shape. It requires questioning assumptions (e.g., not assuming the shape is simple or common) and evaluating the information given in the path data.\\n\\n2. **Creative Thinking (11):** While the task seems straightforward, creative thinking can help in visualizing the shape described by the path commands without immediately drawing it. This involves imagining the transitions and connections between the points defined in the path.\\n\\n3. **Systems Thinking (13):** Understanding the SVG path as a system of coordinates and lines that connect to form a shape. This includes recognizing the interconnectedness of the start and end points of each line segment and how they contribute to the overall shape.\\n\\n4. **Analytical Problem Solving (29):** This task requires data analysis skills to interpret the SVG path commands and deduce the shape they form. Analyzing the coordinates and the movements (lines and moves) can reveal the structure of the shape.\\n\\n5. **Design Challenge (30):** Interpreting and visualizing SVG paths can be seen as a design challenge, requiring an understanding of how individual parts (line segments) come together to create a whole (shape).\\n\\n6. **Step-by-Step Planning and Implementation (39):** Formulating a plan to sequentially interpret each segment of the SVG path and understanding how each segment contributes to the overall shape. This could involve sketching the path based on the commands to better visualize the shape.\\n\\nThese modules collectively enable a comprehensive approach to solving the task, from understanding and analyzing the SVG path data to creatively and systematically deducing the shape it represents.'}}\n",
      "{'adapt': {'adapted_modules': \"To enhance the process of identifying the shape drawn by the SVG path element, the reasoning modules can be adapted and specified as follows:\\n\\n1. **Enhanced Critical Analysis (10):** This module focuses on a detailed examination of the SVG path commands, challenging initial perceptions and critically assessing each command's role in shaping the figure. It involves a deep dive into the syntax and semantics of the path data, ensuring no detail is overlooked, especially in recognizing less obvious or complex shapes.\\n\\n2. **Visual Creative Thinking (11):** Leveraging imagination to mentally construct the shape from the path commands, this module emphasizes the ability to visualize the sequential flow and connection of points without physical drawing. It encourages innovative approaches to mentally piecing together the described shape, enhancing the ability to predict the outcome based on abstract data.\\n\\n3. **Integrated Systems Analysis (13):** This module treats the SVG path as a complex system where each command and coordinate plays a critical role in the final shape. It focuses on understanding the relationship between individual path segments and their collective contribution to forming a coherent structure, emphasizing the holistic view of the path's construction.\\n\\n4. **Targeted Analytical Problem Solving (29):** Specializing in dissecting the SVG path's commands to systematically uncover the represented shape, this module applies precise analytical techniques to decode the sequence of movements and coordinates. It involves a methodical breakdown of the path data to reveal the underlying geometric figure.\\n\\n5. **Design Synthesis Challenge (30):** Approaching the task as a problem of synthesizing a coherent design from segmented inputs, this module requires an adept understanding of how discrete line segments interconnect to form a unified shape. It challenges one to think like a designer, piecing together the puzzle of path commands into a complete and recognizable form.\\n\\n6. **Sequential Interpretation and Visualization (39):** This module involves developing a step-by-step strategy for interpreting and visualizing the SVG path, focusing on the incremental construction of the shape from the path commands. It advocates for a systematic approach to translating the abstract commands into a tangible visual representation, potentially through sketching or mentally mapping the path's progression.\\n\\nBy refining these modules, the approach to solving the task becomes more targeted, enhancing the ability to accurately identify the shape described by the SVG path element.\"}}\n"
     ]
    }
   ],
   "source": [
    "reasoning_modules = [\n",
    "    \"1. How could I devise an experiment to help solve that problem?\",\n",
    "    \"2. Make a list of ideas for solving this problem, and apply them one by one to the problem to see if any progress can be made.\",\n",
    "    # \"3. How could I measure progress on this problem?\",\n",
    "    \"4. How can I simplify the problem so that it is easier to solve?\",\n",
    "    \"5. What are the key assumptions underlying this problem?\",\n",
    "    \"6. What are the potential risks and drawbacks of each solution?\",\n",
    "    \"7. What are the alternative perspectives or viewpoints on this problem?\",\n",
    "    \"8. What are the long-term implications of this problem and its solutions?\",\n",
    "    \"9. How can I break down this problem into smaller, more manageable parts?\",\n",
    "    \"10. Critical Thinking: This style involves analyzing the problem from different perspectives, questioning assumptions, and evaluating the evidence or information available. It focuses on logical reasoning, evidence-based decision-making, and identifying potential biases or flaws in thinking.\",\n",
    "    \"11. Try creative thinking, generate innovative and out-of-the-box ideas to solve the problem. Explore unconventional solutions, thinking beyond traditional boundaries, and encouraging imagination and originality.\",\n",
    "    # \"12. Seek input and collaboration from others to solve the problem. Emphasize teamwork, open communication, and leveraging the diverse perspectives and expertise of a group to come up with effective solutions.\",\n",
    "    \"13. Use systems thinking: Consider the problem as part of a larger system and understanding the interconnectedness of various elements. Focuses on identifying the underlying causes, feedback loops, and interdependencies that influence the problem, and developing holistic solutions that address the system as a whole.\",\n",
    "    \"14. Use Risk Analysis: Evaluate potential risks, uncertainties, and tradeoffs associated with different solutions or approaches to a problem. Emphasize assessing the potential consequences and likelihood of success or failure, and making informed decisions based on a balanced analysis of risks and benefits.\",\n",
    "    # \"15. Use Reflective Thinking: Step back from the problem, take the time for introspection and self-reflection. Examine personal biases, assumptions, and mental models that may influence problem-solving, and being open to learning from past experiences to improve future approaches.\",\n",
    "    \"16. What is the core issue or problem that needs to be addressed?\",\n",
    "    \"17. What are the underlying causes or factors contributing to the problem?\",\n",
    "    \"18. Are there any potential solutions or strategies that have been tried before? If yes, what were the outcomes and lessons learned?\",\n",
    "    \"19. What are the potential obstacles or challenges that might arise in solving this problem?\",\n",
    "    \"20. Are there any relevant data or information that can provide insights into the problem? If yes, what data sources are available, and how can they be analyzed?\",\n",
    "    \"21. Are there any stakeholders or individuals who are directly affected by the problem? What are their perspectives and needs?\",\n",
    "    \"22. What resources (financial, human, technological, etc.) are needed to tackle the problem effectively?\",\n",
    "    \"23. How can progress or success in solving the problem be measured or evaluated?\",\n",
    "    \"24. What indicators or metrics can be used?\",\n",
    "    \"25. Is the problem a technical or practical one that requires a specific expertise or skill set? Or is it more of a conceptual or theoretical problem?\",\n",
    "    \"26. Does the problem involve a physical constraint, such as limited resources, infrastructure, or space?\",\n",
    "    \"27. Is the problem related to human behavior, such as a social, cultural, or psychological issue?\",\n",
    "    \"28. Does the problem involve decision-making or planning, where choices need to be made under uncertainty or with competing objectives?\",\n",
    "    \"29. Is the problem an analytical one that requires data analysis, modeling, or optimization techniques?\",\n",
    "    \"30. Is the problem a design challenge that requires creative solutions and innovation?\",\n",
    "    \"31. Does the problem require addressing systemic or structural issues rather than just individual instances?\",\n",
    "    \"32. Is the problem time-sensitive or urgent, requiring immediate attention and action?\",\n",
    "    \"33. What kinds of solution typically are produced for this kind of problem specification?\",\n",
    "    \"34. Given the problem specification and the current best solution, have a guess about other possible solutions.\"\n",
    "    \"35. Let’s imagine the current best solution is totally wrong, what other ways are there to think about the problem specification?\"\n",
    "    \"36. What is the best way to modify this current best solution, given what you know about these kinds of problem specification?\"\n",
    "    \"37. Ignoring the current best solution, create an entirely new solution to the problem.\"\n",
    "    # \"38. Let’s think step by step.\"\n",
    "    \"39. Let’s make a step by step plan and implement it with good notation and explanation.\",\n",
    "]\n",
    "\n",
    "\n",
    "task_example = \"Lisa has 10 apples. She gives 3 apples to her friend and then buys 5 more apples from the store. How many apples does Lisa have now?\"\n",
    "\n",
    "task_example = \"\"\"This SVG path element <path d=\"M 55.57,80.69 L 57.38,65.80 M 57.38,65.80 L 48.90,57.46 M 48.90,57.46 L\n",
    "45.58,47.78 M 45.58,47.78 L 53.25,36.07 L 66.29,48.90 L 78.69,61.09 L 55.57,80.69\"/> draws a:\n",
    "(A) circle (B) heptagon (C) hexagon (D) kite (E) line (F) octagon (G) pentagon(H) rectangle (I) sector (J) triangle\"\"\"\n",
    "\n",
    "reasoning_modules_str = \"\\n\".join(reasoning_modules)\n",
    "\n",
    "for s in app.stream(\n",
    "    {\"task_description\": task_example, \"reasoning_modules\": reasoning_modules_str}\n",
    "):\n",
    "    print(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20cac598",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
