{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d72fae4e-f7de-42b7-91ee-bdd0a57ae46c",
   "metadata": {},
   "source": [
    "# Prompt Generation from User Requirements\n",
    "\n",
    "In this example we will create a chat bot that helps a user generate a prompt.\n",
    "It will first collect requirements from the user, and then will generate the prompt (and refine it based on user input).\n",
    "These are split into two separate states, and the LLM decides when to transition between them.\n",
    "\n",
    "A graphical representation of the system can be found below.\n",
    "\n",
    "![](imgs/prompt-generator.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d78b593-ba26-4c90-b2e2-83119e47679f",
   "metadata": {},
   "source": [
    "## Gather information\n",
    "\n",
    "First, let's define the part of the graph that will gather user requirements. This will be an LLM call with a specific system message. It will have access to a tool that it can call when it is ready to generate the prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "53216ab5-2cd3-48a4-8778-41ba10f72519",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "from langchain_core.messages import SystemMessage\n",
    "from langchain_core.pydantic_v1 import BaseModel\n",
    "from langchain_openai import ChatOpenAI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5f795b78-004d-40ca-95d6-069f67e4f9c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "template = \"\"\"Your job is to get information from a user about what type of prompt template they want to create.\n",
    "\n",
    "You should get the following information from them:\n",
    "\n",
    "- What the objective of the prompt is\n",
    "- What variables will be passed into the prompt template\n",
    "- Any constraints for what the output should NOT do\n",
    "- Any requirements that the output MUST adhere to\n",
    "\n",
    "If you are not able to discern this info, ask them to clarify! Do not attempt to wildly guess.\n",
    "\n",
    "After you are able to discern all the information, call the relevant tool.\"\"\"\n",
    "\n",
    "\n",
    "def get_messages_info(messages):\n",
    "    return [SystemMessage(content=template)] + messages\n",
    "\n",
    "\n",
    "class PromptInstructions(BaseModel):\n",
    "    \"\"\"Instructions on how to prompt the LLM.\"\"\"\n",
    "\n",
    "    objective: str\n",
    "    variables: List[str]\n",
    "    constraints: List[str]\n",
    "    requirements: List[str]\n",
    "\n",
    "\n",
    "llm = ChatOpenAI(temperature=0)\n",
    "llm_with_tool = llm.bind_tools([PromptInstructions])\n",
    "\n",
    "chain = get_messages_info | llm_with_tool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb40630f-83c7-4283-a6dd-04231805a7ed",
   "metadata": {},
   "source": [
    "## Generate Prompt\n",
    "\n",
    "We now set up the state that will generate the prompt.\n",
    "This will require a separate system message, as well as a function to filter out all message PRIOR to the tool invocation (as that is when the previous state decided it was time to generate the prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca9a0234-bbeb-4bff-8276-8dde499c3390",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import AIMessage, HumanMessage, ToolMessage\n",
    "\n",
    "# New system prompt\n",
    "prompt_system = \"\"\"Based on the following requirements, write a good prompt template:\n",
    "\n",
    "{reqs}\"\"\"\n",
    "\n",
    "\n",
    "# Function to get the messages for the prompt\n",
    "# Will only get messages AFTER the tool call\n",
    "def get_prompt_messages(messages: list):\n",
    "    tool_call = None\n",
    "    other_msgs = []\n",
    "    for m in messages:\n",
    "        if isinstance(m, AIMessage) and m.tool_calls:\n",
    "            tool_call = m.tool_calls[0][\"args\"]\n",
    "        elif isinstance(m, ToolMessage):\n",
    "            continue\n",
    "        elif tool_call is not None:\n",
    "            other_msgs.append(m)\n",
    "    return [SystemMessage(content=prompt_system.format(reqs=tool_call))] + other_msgs\n",
    "\n",
    "\n",
    "prompt_gen_chain = get_prompt_messages | llm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dbabda8-34f0-4eef-bce2-ad3ff505366b",
   "metadata": {},
   "source": [
    "## Define the state logic\n",
    "\n",
    "This is the logic for what state the chatbot is in.\n",
    "If the last message is a tool call, then we are in the state where the \"prompt creator\" (`prompt`) should respond.\n",
    "Otherwise, if the last message is not a HumanMessage, then we know the human should respond next and so we are in the `END` state.\n",
    "If the last message is a HumanMessage, then if there was a tool call previously we are in the `prompt` state.\n",
    "Otherwise, we are in the \"info gathering\" (`info`) state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f29e15-20e2-420c-a450-84e929f16e4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "\n",
    "from langgraph.graph import END\n",
    "\n",
    "\n",
    "def get_state(messages) -> Literal[\"add_tool_message\", \"info\", \"__end__\"]:\n",
    "    if isinstance(messages[-1], AIMessage) and messages[-1].tool_calls:\n",
    "        return \"add_tool_message\"\n",
    "    elif not isinstance(messages[-1], HumanMessage):\n",
    "        return END\n",
    "    return \"info\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b76bea78-07a5-418f-9b7c-71c376d4b6f7",
   "metadata": {},
   "source": [
    "## Create the graph\n",
    "\n",
    "We can now the create the graph.\n",
    "We will use a SqliteSaver to persist conversation history."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59d9d6b4-dce4-43cc-9a1a-61a7912ed5b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import StateGraph, START\n",
    "from langgraph.graph.message import add_messages\n",
    "from typing import Annotated\n",
    "from typing_extensions import TypedDict\n",
    "\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[list, add_messages]\n",
    "\n",
    "memory = MemorySaver()\n",
    "workflow = StateGraph(State)\n",
    "workflow.add_node(\"info\", chain)\n",
    "workflow.add_node(\"prompt\", prompt_gen_chain)\n",
    "\n",
    "\n",
    "@workflow.add_node\n",
    "def add_tool_message(state: list):\n",
    "    return ToolMessage(\n",
    "        content=\"Prompt generated!\", tool_call_id=state[-1].tool_calls[0][\"id\"]\n",
    "    )\n",
    "\n",
    "\n",
    "workflow.add_conditional_edges(\"info\", get_state)\n",
    "workflow.add_edge(\"add_tool_message\", \"prompt\")\n",
    "workflow.add_edge(\"prompt\", END)\n",
    "workflow.add_edge(START, \"info\")\n",
    "graph = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "1b1613e0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 17416 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import Image, display\n",
    "\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afcf523c-265d-45cf-a981-fc50c50c1738",
   "metadata": {},
   "source": [
    "## Use the graph\n",
    "\n",
    "We can now use the created chatbot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "25793988-45a2-4e65-b33c-64e72aadb10e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Hello! How can I assist you today?\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Sure! I can help you with that. To create an extraction prompt, I need some information from you. Could you please provide the following details:\n",
      "\n",
      "1. What is the objective of the prompt?\n",
      "2. What variables will be passed into the prompt template?\n",
      "3. Any constraints for what the output should NOT do?\n",
      "4. Any requirements that the output MUST adhere to?\n",
      "\n",
      "Once I have this information, I can create the extraction prompt for you.\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Great! To create an extraction prompt for filling out a CSAT (Customer Satisfaction) survey, I will need the following information:\n",
      "\n",
      "1. Objective: To gather feedback on customer satisfaction.\n",
      "2. Variables: Customer name, Date of interaction, Service provided, Rating (scale of 1-5), Comments.\n",
      "3. Constraints: The output should not include any personally identifiable information (PII) of the customer.\n",
      "4. Requirements: The output must include a structured format with fields for each variable mentioned above.\n",
      "\n",
      "With this information, I will proceed to create the extraction prompt template for filling out a CSAT survey. Let's get started!\n",
      "Tool Calls:\n",
      "  PromptInstructions (call_aU48Bjo7X29tXfRtCcrXkrqq)\n",
      " Call ID: call_aU48Bjo7X29tXfRtCcrXkrqq\n",
      "  Args:\n",
      "    objective: To gather feedback on customer satisfaction.\n",
      "    variables: ['Customer name', 'Date of interaction', 'Service provided', 'Rating (scale of 1-5)', 'Comments']\n",
      "    constraints: ['The output should not include any personally identifiable information (PII) of the customer.']\n",
      "    requirements: ['The output must include a structured format with fields for each variable mentioned above.']\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "\n",
      "Prompt generated!\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Please provide feedback on your recent interaction with our service. Your input is valuable to us in improving our services.\n",
      "\n",
      "Customer name: \n",
      "Date of interaction: \n",
      "Service provided: \n",
      "Rating (scale of 1-5): \n",
      "Comments: \n",
      "\n",
      "Please note that the output should not include any personally identifiable information (PII) of the customer. Your feedback will be kept confidential and used for internal evaluation purposes only. Thank you for taking the time to share your thoughts with us.\n",
      "Done!\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "I'm glad you found it helpful! If you need any more assistance or have any other requests, feel free to let me know. Have a great day!\n",
      "AI: Byebye\n"
     ]
    }
   ],
   "source": [
    "import uuid\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": str(uuid.uuid4())}}\n",
    "while True:\n",
    "    user = input(\"User (q/Q to quit): \")\n",
    "    if user in {\"q\", \"Q\"}:\n",
    "        print(\"AI: Byebye\")\n",
    "        break\n",
    "    output = None\n",
    "    for output in graph.stream(\n",
    "        [HumanMessage(content=user)], config=config, stream_mode=\"updates\"\n",
    "    ):\n",
    "        last_message = next(iter(output.values()))\n",
    "        last_message.pretty_print()\n",
    "\n",
    "    if output and \"prompt\" in output:\n",
    "        print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a276d20e-8a1b-4add-bf8d-83a8c803431d",
   "metadata": {},
   "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
}
