examples/agents/langgraph_javelin.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting langgraph\n",
      "  Using cached langgraph-0.2.69-py3-none-any.whl.metadata (17 kB)\n",
      "Requirement already satisfied: langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langgraph) (0.3.32)\n",
      "Collecting langgraph-checkpoint<3.0.0,>=2.0.10 (from langgraph)\n",
      "  Using cached langgraph_checkpoint-2.0.10-py3-none-any.whl.metadata (4.6 kB)\n",
      "Collecting langgraph-sdk<0.2.0,>=0.1.42 (from langgraph)\n",
      "  Using cached langgraph_sdk-0.1.51-py3-none-any.whl.metadata (1.8 kB)\n",
      "Requirement already satisfied: PyYAML>=5.3 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (6.0.2)\n",
      "Requirement already satisfied: jsonpatch<2.0,>=1.33 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (1.33)\n",
      "Requirement already satisfied: langsmith<0.4,>=0.1.125 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (0.3.2)\n",
      "Requirement already satisfied: packaging<25,>=23.2 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (24.2)\n",
      "Requirement already satisfied: pydantic<3.0.0,>=2.7.4 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (2.10.6)\n",
      "Requirement already satisfied: tenacity!=8.4.0,<10.0.0,>=8.1.0 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (9.0.0)\n",
      "Requirement already satisfied: typing-extensions>=4.7 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (4.12.2)\n",
      "Collecting msgpack<2.0.0,>=1.1.0 (from langgraph-checkpoint<3.0.0,>=2.0.10->langgraph)\n",
      "  Using cached msgpack-1.1.0-cp312-cp312-macosx_11_0_arm64.whl.metadata (8.4 kB)\n",
      "Collecting httpx>=0.25.2 (from langgraph-sdk<0.2.0,>=0.1.42->langgraph)\n",
      "  Using cached httpx-0.28.1-py3-none-any.whl.metadata (7.1 kB)\n",
      "Requirement already satisfied: orjson>=3.10.1 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langgraph-sdk<0.2.0,>=0.1.42->langgraph) (3.10.15)\n",
      "Requirement already satisfied: anyio in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from httpx>=0.25.2->langgraph-sdk<0.2.0,>=0.1.42->langgraph) (4.8.0)\n",
      "Requirement already satisfied: certifi in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from httpx>=0.25.2->langgraph-sdk<0.2.0,>=0.1.42->langgraph) (2024.12.14)\n",
      "Collecting httpcore==1.* (from httpx>=0.25.2->langgraph-sdk<0.2.0,>=0.1.42->langgraph)\n",
      "  Using cached httpcore-1.0.7-py3-none-any.whl.metadata (21 kB)\n",
      "Requirement already satisfied: idna in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from httpx>=0.25.2->langgraph-sdk<0.2.0,>=0.1.42->langgraph) (3.10)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from httpcore==1.*->httpx>=0.25.2->langgraph-sdk<0.2.0,>=0.1.42->langgraph) (0.14.0)\n",
      "Requirement already satisfied: jsonpointer>=1.9 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from jsonpatch<2.0,>=1.33->langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (3.0.0)\n",
      "Requirement already satisfied: requests<3,>=2 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langsmith<0.4,>=0.1.125->langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (2.32.3)\n",
      "Requirement already satisfied: requests-toolbelt<2.0.0,>=1.0.0 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langsmith<0.4,>=0.1.125->langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (1.0.0)\n",
      "Requirement already satisfied: zstandard<0.24.0,>=0.23.0 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from langsmith<0.4,>=0.1.125->langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (0.23.0)\n",
      "Requirement already satisfied: annotated-types>=0.6.0 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.7.4->langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (0.7.0)\n",
      "Requirement already satisfied: pydantic-core==2.27.2 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from pydantic<3.0.0,>=2.7.4->langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (2.27.2)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from requests<3,>=2->langsmith<0.4,>=0.1.125->langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (3.4.1)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from requests<3,>=2->langsmith<0.4,>=0.1.125->langchain-core!=0.3.0,!=0.3.1,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14,!=0.3.15,!=0.3.16,!=0.3.17,!=0.3.18,!=0.3.19,!=0.3.2,!=0.3.20,!=0.3.21,!=0.3.22,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,<0.4.0,>=0.2.43->langgraph) (2.3.0)\n",
      "Requirement already satisfied: sniffio>=1.1 in /Users/dhruvyadav/Desktop/javelin-main/javelin-python/venv/lib/python3.12/site-packages (from anyio->httpx>=0.25.2->langgraph-sdk<0.2.0,>=0.1.42->langgraph) (1.3.1)\n",
      "Using cached langgraph-0.2.69-py3-none-any.whl (148 kB)\n",
      "Using cached langgraph_checkpoint-2.0.10-py3-none-any.whl (37 kB)\n",
      "Using cached langgraph_sdk-0.1.51-py3-none-any.whl (44 kB)\n",
      "Using cached httpx-0.28.1-py3-none-any.whl (73 kB)\n",
      "Using cached httpcore-1.0.7-py3-none-any.whl (78 kB)\n",
      "Using cached msgpack-1.1.0-cp312-cp312-macosx_11_0_arm64.whl (82 kB)\n",
      "Installing collected packages: msgpack, httpcore, httpx, langgraph-sdk, langgraph-checkpoint, langgraph\n",
      "  Attempting uninstall: httpcore\n",
      "    Found existing installation: httpcore 0.17.3\n",
      "    Uninstalling httpcore-0.17.3:\n",
      "      Successfully uninstalled httpcore-0.17.3\n",
      "  Attempting uninstall: httpx\n",
      "    Found existing installation: httpx 0.24.1\n",
      "    Uninstalling httpx-0.24.1:\n",
      "      Successfully uninstalled httpx-0.24.1\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "javelin-sdk 0.2.19 requires httpx<0.25.0,>=0.24.0, but you have httpx 0.28.1 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0mSuccessfully installed httpcore-1.0.7 httpx-0.28.1 langgraph-0.2.69 langgraph-checkpoint-2.0.10 langgraph-sdk-0.1.51 msgpack-1.1.0\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "pip install langgraph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Email Generation Workflow using Javelin and LangGraph\n",
    "\n",
    "## Overview\n",
    "\n",
    "This script uses the **Javelin API** and **LangGraph** to perform email generation based on user queries. The flow involves validating the query, refining it, and generating the email.\n",
    "\n",
    "## Key Components\n",
    "\n",
    "1. **Javelin API**: \n",
    "   - The Javelin API is used to validate and refine the user's query by providing a route (`testing` in this case). It helps assess whether the request is suitable for email generation and aids in refining the query for clarity.\n",
    "   - **Headers**: Contains the API key and the route for Javelin (`x-javelin-route`).\n",
    "   \n",
    "2. **LangGraph**:\n",
    "   - LangGraph is used to manage the flow of agents, where each agent is responsible for a specific task: \n",
    "     - **Agent 1**: Validates if the query can be turned into an email.\n",
    "     - **Agent 2**: Refines the query.\n",
    "     - **Agent 3**: Generates the email from the refined query.\n",
    "   \n",
    "3. **StateGraph**:\n",
    "   - LangGraph’s `StateGraph` is used to create a flowchart-like structure for the process. It connects the agents in a sequence, ensuring that the right steps are followed in the correct order.\n",
    "\n",
    "## Workflow\n",
    "\n",
    "- **Step 1**: The user's query is first passed to **Agent 1**, where it's validated by querying Javelin. If valid, it moves to **Agent 2**.\n",
    "- **Step 2**: **Agent 2** refines the query for clarity, ensuring that the purpose and recipients are clear.\n",
    "- **Step 3**: **Agent 3** takes the refined query and generates an email.\n",
    "- The system continues its flow until the final email is generated.\n",
    "\n",
    "## How Javelin Plays a Role\n",
    "\n",
    "Javelin is used as a \"router\" for query validation and refinement:\n",
    "- It provides the necessary decision-making process by analyzing the query and returning whether it's suitable for email generation.\n",
    "- The API returns a JSON response that determines the flow (valid/invalid).\n",
    "\n",
    "## LangGraph’s Role\n",
    "\n",
    "LangGraph is responsible for managing the execution flow:\n",
    "- It ensures the agents perform their tasks in the correct order.\n",
    "- It allows for dynamic branching, ensuring that if the query is not valid, the process stops.\n",
    "\n",
    "## Final Goal\n",
    "\n",
    "The final goal is to:\n",
    "1. Validate the user’s query.\n",
    "2. Refine it if necessary.\n",
    "3. Generate a valid email based on the refined query.\n",
    "\n",
    "Each of these steps is achieved through the seamless interaction between Javelin (for validation) and LangGraph (for orchestrating the flow of tasks).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "User Query: I want to send an email to my boss to request more leave because I am feeling unwell.\n",
      "Final Response: Subject: Request for Additional Leave Due to Illness\n",
      "\n",
      "Dear [Boss's Name],\n",
      "\n",
      "I am writing to request additional leave from work due to feeling unwell. I have been experiencing worsening symptoms that have made it difficult to fully recover within the initially scheduled time off. I am reaching out to seek your understanding and support during this time.\n",
      "\n",
      "I believe that taking additional time off to focus on my health will allow me to recuperate fully and return to work with renewed energy and focus. I have been in touch with my healthcare provider and am following their guidance to ensure a swift recovery.\n",
      "\n",
      "I understand the importance of my role and the impact of my absence on the team. I am committed to staying updated on any urgent matters and ensuring a smooth transition of responsibilities during my extended leave.\n",
      "\n",
      "I would appreciate your guidance on the process for requesting additional leave and any necessary documentation that may be required. Please let me know if there are any specific procedures I should follow or if there are alternative arrangements that need to be made to cover my responsibilities during my absence.\n",
      "\n",
      "Thank you for your understanding and support. I look forward to your feedback and advice regarding my request for additional leave.\n",
      "\n",
      "Warm regards,\n",
      "\n",
      "[Your Name]\n",
      "\n",
      "User Query: Who is the President of India?\n",
      "Final Response: No response.\n"
     ]
    }
   ],
   "source": [
    "from openai import OpenAI\n",
    "import os\n",
    "import json\n",
    "from typing import List, TypedDict, Literal\n",
    "from langgraph.graph import StateGraph, END\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "# -------------------------------\n",
    "# Initialize Clients using Unified Endpoint\n",
    "# -------------------------------\n",
    "# Set your API keys in your environment variables: OPENAI_API_KEY and JAVELIN_API_KEY\n",
    "openai_api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "javelin_api_key = os.getenv(\"JAVELIN_API_KEY\")\n",
    "\n",
    "# Create a plain OpenAI client\n",
    "openai_client = OpenAI(api_key=openai_api_key)\n",
    "\n",
    "# Initialize Javelin unified endpoint client\n",
    "from javelin_sdk import JavelinClient, JavelinConfig\n",
    "\n",
    "config = JavelinConfig(javelin_api_key=javelin_api_key)\n",
    "client = JavelinClient(config)\n",
    "# Register the OpenAI client with the unified route name (e.g., \"openai_univ\")\n",
    "client.register_openai(openai_client, route_name=\"openai_univ\")\n",
    "\n",
    "# -------------------------------\n",
    "# Define Message Structure\n",
    "# -------------------------------\n",
    "class MessagesState(TypedDict):\n",
    "    messages: List[dict]\n",
    "    valid: bool            # Indicates if it's a valid email request\n",
    "    refined_query: str     # Stores the refined query (if applicable)\n",
    "\n",
    "# -------------------------------\n",
    "# Agent Functions Using Unified Endpoint\n",
    "# -------------------------------\n",
    "def agent_1(state: MessagesState) -> MessagesState:\n",
    "    messages = state[\"messages\"]\n",
    "    user_message = messages[-1][\"content\"].lower()\n",
    "\n",
    "    # Validate if the query is valid for email generation\n",
    "    completion = openai_client.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "            {\"role\": \"user\", \"content\": (\n",
    "                f\"Analyze the following user request: '{user_message}'. \"\n",
    "                \"Determine if it's valid for generating an email based on the request. \"\n",
    "                \"The request is valid if it specifies who to write to, why the email is being written, \"\n",
    "                \"and what to include in the email. If the request does not pertain to generating an email, \"\n",
    "                \"return false for validity. Please justify the validity status in the response. \"\n",
    "                \"Return a JSON object with 'valid' (true/false), 'response', and 'extracted_info' \"\n",
    "                \"with details such as recipient and reason.\"\n",
    "            )}\n",
    "        ]\n",
    "    )\n",
    "    \n",
    "    response_json = completion.choices[0].message.content.strip()\n",
    "    try:\n",
    "        parsed_response = json.loads(response_json)\n",
    "        valid = parsed_response.get(\"valid\", False)\n",
    "    except json.JSONDecodeError:\n",
    "        valid = False\n",
    "\n",
    "    return {\n",
    "        \"messages\": messages,\n",
    "        \"valid\": valid,\n",
    "        \"refined_query\": \"\"\n",
    "    }\n",
    "\n",
    "# Agent 2: Refine the query if valid\n",
    "def agent_2(state: MessagesState) -> MessagesState:\n",
    "    if not state[\"valid\"]:\n",
    "        return state  # Skip if invalid\n",
    "\n",
    "    messages = state[\"messages\"]\n",
    "    user_message = messages[-1][\"content\"]\n",
    "\n",
    "    # Refine the query if valid\n",
    "    completion = openai_client.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": (\n",
    "                \"You are a helpful assistant to refine the query and highlight the points.\"\n",
    "            )},\n",
    "            {\"role\": \"user\", \"content\": (\n",
    "                f\"Refine the following query for email generation: '{user_message}'. \"\n",
    "                \"Ensure to highlight who to send the email to, why, and what needs to be included. \"\n",
    "                \"Important: do not create an email. Just return a refined user query.\"\n",
    "            )}\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    refined_query = completion.choices[0].message.content.strip()\n",
    "\n",
    "    return {\n",
    "        \"messages\": messages,\n",
    "        \"valid\": True,\n",
    "        \"refined_query\": refined_query\n",
    "    }\n",
    "\n",
    "def agent_3(state: MessagesState) -> MessagesState:\n",
    "    if not state[\"refined_query\"]:\n",
    "        return state\n",
    "\n",
    "    refined_query = state[\"refined_query\"]\n",
    "\n",
    "    # Generate the email based on the refined query\n",
    "    completion = openai_client.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "            {\"role\": \"user\", \"content\": f\"Generate an email based on the following details: '{refined_query}'\"}\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    email_content = completion.choices[0].message.content.strip()\n",
    "\n",
    "    return {\n",
    "        \"messages\": state[\"messages\"] + [{\"role\": \"assistant\", \"content\": email_content}],\n",
    "        \"valid\": True,\n",
    "        \"refined_query\": refined_query\n",
    "    }\n",
    "\n",
    "# Decision function to determine flow\n",
    "def should_continue(state: MessagesState) -> Literal[\"tools\", END]:\n",
    "    if state[\"valid\"]:\n",
    "        if not state[\"refined_query\"]:\n",
    "            print(\"should_continue decision: tools (Proceeding to query refinement)\")\n",
    "            return \"tools\"\n",
    "        else:\n",
    "            print(\"should_continue decision: tools (Proceeding to email generation)\")\n",
    "            return \"tools\"\n",
    "    else:\n",
    "        print(\"should_continue decision: __end__ (Stopping, invalid request)\")\n",
    "        return END\n",
    "\n",
    "# -------------------------------\n",
    "# Set Up the State Machine\n",
    "# -------------------------------\n",
    "graph = StateGraph(MessagesState)\n",
    "\n",
    "# Define the nodes for each agent\n",
    "graph.add_node(\"agent_1\", agent_1)\n",
    "graph.add_node(\"agent_2\", agent_2)\n",
    "graph.add_node(\"agent_3\", agent_3)\n",
    "\n",
    "# Define the flow of agents\n",
    "graph.add_edge(\"agent_1\", \"agent_2\")  # Proceed to agent 2 if agent 1 validates\n",
    "graph.add_edge(\"agent_2\", \"agent_3\")  # Proceed to agent 3 if agent 2 refines the query\n",
    "\n",
    "# Set the entry point for the state machine\n",
    "graph.set_entry_point(\"agent_1\")\n",
    "\n",
    "# Compile the graph\n",
    "app = graph.compile()\n",
    "\n",
    "# Test cases\n",
    "queries = [\n",
    "    \"I want to send an email to my boss to request more leave because I am feeling unwell.\",\n",
    "    \"Who is the President of India?\"\n",
    "]\n",
    "\n",
    "for query in queries:\n",
    "    print(\"\\nUser Query:\", query)\n",
    "    initial_state = {\n",
    "        \"messages\": [{\"role\": \"user\", \"content\": query}],\n",
    "        \"valid\": True,\n",
    "        \"refined_query\": \"\"\n",
    "    }\n",
    "\n",
    "    final_state = app.invoke(initial_state, config={\"debug\": True})\n",
    "    assistant_response = next(\n",
    "        (msg[\"content\"] for msg in final_state[\"messages\"] if msg[\"role\"] == \"assistant\"),\n",
    "        \"No response.\"\n",
    "    )\n",
    "    print(\"Final Response:\", assistant_response)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}