contextual_quoting_agentic_system.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# **Dynamic Quoting System for Complex Products Using Multi-Agent AI Architecture**\n",
    "\n",
    "## **Project Overview**\n",
    "This project implements a sophisticated multi-agent system for generating personalized, context-aware quotes for complex products and services. It leverages LangChain/LangGraph to create an intelligent quoting workflow that goes beyond traditional pricing methods by incorporating:\n",
    "\n",
    "- **Retrieval-Augmented Generation (RAG)** for context-sensitive data retrieval\n",
    "- **Multi-Agent Architecture** with specialized roles:\n",
    "  - **Main Assistant**: Initial information gathering\n",
    "  - **Underwriting Assistant**: Risk evaluation\n",
    "  - **Quote Assistant**: Premium calculation\n",
    "- **Intelligent Classification** using sentiment analysis and business context\n",
    "- **Dynamic Workflow Management** through a state graph system\n",
    "- **Database Integration** for storing and retrieving category rates\n",
    "\n",
    "## **Motivation**\n",
    "Traditional quoting systems often struggle with complex products and services where pricing depends on multiple interrelated factors. Current solutions typically:\n",
    "- Rely heavily on manual intervention\n",
    "- Have limited ability to consider context\n",
    "- Struggle with non-standard cases\n",
    "- Lack consistency across different underwriters\n",
    "- Cannot easily adapt to changing market conditions\n",
    "\n",
    "This project addresses these limitations by creating an intelligent system that can understand context, learn from historical data, and provide consistent, accurate quotes while reducing manual effort and turnaround time.\n",
    "\n",
    "## **General Use Cases**\n",
    "This system is particularly valuable in industries such as:\n",
    "\n",
    "1. **Insurance**\n",
    "   - Commercial insurance quoting\n",
    "   - Risk assessment and premium calculation\n",
    "   - Policy customization based on business specifics\n",
    "\n",
    "2. **Professional Services**\n",
    "   - Consulting service pricing\n",
    "   - Project cost estimation\n",
    "   - Service package customization\n",
    "\n",
    "3. **Manufacturing**\n",
    "   - Custom product pricing\n",
    "   - Bill of materials calculation\n",
    "   - Production cost estimation\n",
    "\n",
    "4. **Construction**\n",
    "   - Project bidding\n",
    "   - Material and labor cost estimation\n",
    "   - Risk-adjusted pricing\n",
    "\n",
    "## **Key Components**\n",
    "\n",
    "1. **Core Infrastructure**\n",
    "   - SQLite database for storing category rates\n",
    "   - Pydantic schemas for data validation\n",
    "   - State management using TypedDict\n",
    "   - LangGraph for workflow orchestration\n",
    "\n",
    "2. **Specialized Agents**\n",
    "   - **Retriever Agent**: Extracts and summarizes business operations\n",
    "   - **Reasoning Agent**: Determines relevant insurance categories\n",
    "   - **Classification Grading Agent**: Evaluates category assignments\n",
    "   - **Quote Generation Agent**: Calculates final premiums\n",
    "\n",
    "3. **Workflow Management**\n",
    "   - Dynamic routing between agents\n",
    "   - State tracking across conversation flows\n",
    "   - Error handling and fallback mechanisms\n",
    "   - Tool management for specific actions\n",
    "\n",
    "## **Technical Implementation**\n",
    "The system uses a graph-based architecture where nodes represent different stages of the quoting process and edges define the possible transitions between these stages. The workflow is managed through:\n",
    "\n",
    "- **State Management**: Tracking conversation context and business information\n",
    "- **Routing Logic**: Determining which agent handles each step\n",
    "- **Tool Integration**: Specialized functions for specific tasks\n",
    "- **Memory Management**: Maintaining conversation history and context\n",
    "\n",
    "## **Conclusion**\n",
    "This multi-agent quoting system represents a significant advancement in automated pricing solutions. By combining AI-driven analysis, context awareness, and flexible workflow management, it addresses the complexities of modern pricing challenges while providing several key benefits:\n",
    "\n",
    "### **Business Impact**\n",
    "- **Reduced Manual Effort**: Automation of complex pricing decisions\n",
    "- **Increased Consistency**: Standardized approach across all quotes\n",
    "- **Faster Turnaround**: Reduced time from inquiry to quote\n",
    "- **Better Accuracy**: Context-aware pricing decisions\n",
    "- **Scalability**: Ability to handle increased quote volumes without proportional staff increases\n",
    "\n",
    "### **Technical Achievement**\n",
    "- Demonstrates the power of multi-agent architectures in solving complex business problems\n",
    "- Shows how different AI components can work together in a coordinated workflow\n",
    "- Provides a blueprint for building similar systems in other domains\n",
    "\n",
    "### **Future Potential**\n",
    "The architecture can be extended to include:\n",
    "- Additional specialized agents for specific industries\n",
    "- More sophisticated pricing models\n",
    "- Integration with external data sources\n",
    "- Machine learning components for continuous improvement\n",
    "\n",
    "This system serves as a model for how AI can transform complex business processes, making them more efficient, accurate, and scalable while maintaining the flexibility to adapt to changing business needs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import Necessary Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! pip install langgraph\n",
    "# ! pip install langchain-core \n",
    "# ! pip install langchain-openai \n",
    "# ! pip install langchain-groq \n",
    "# ! pip install langchain-community \n",
    "# ! pip install python-dotenv \n",
    "# ! pip install pydantic \n",
    "# ! pip install typing-extensions \n",
    "# ! pip install chromadb \n",
    "# ! pip install langchain-text-splitters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict, TypedDict\n",
    "from langgraph.graph import StateGraph, END\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_groq import ChatGroq\n",
    "from IPython.display import display, Image\n",
    "from langchain_core.runnables.graph import MermaidDrawMethod\n",
    "from dotenv import load_dotenv\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Variables and Enable Tracing for LangGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, getpass\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "_set_env(\"OPENAI_API_KEY\")\n",
    "_set_env(\"LANGCHAIN_API_KEY\")\n",
    "_set_env(\"GROQ_API_KEY\")\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"Hackathon-Contextual-Quoting\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LLM "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To start with, we define the LLM that powers our contextual quoting agents. Here’s the code for initializing the LLM:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "# llm = ChatGroq(model=\"llama3-groq-8b-8192-tool-use-preview\", temperature=0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create and Populate Database"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we create and populate a database to store category rates for our contextual quoting agent. \n",
    "\n",
    "We use **SQLite**, a lightweight and efficient database, which works well for projects where simplicity and speed are key considerations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 200,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Database and table created successfully, and data inserted.\n"
     ]
    }
   ],
   "source": [
    "import sqlite3\n",
    "\n",
    "# Connect to SQLite (or create the database if it doesn't exist)\n",
    "conn = sqlite3.connect('data/categories.db')\n",
    "\n",
    "# Create a cursor object to execute SQL commands\n",
    "cursor = conn.cursor()\n",
    "\n",
    "# Create the category_rates table\n",
    "cursor.execute('''\n",
    "    CREATE TABLE IF NOT EXISTS category_rates (\n",
    "        id INTEGER PRIMARY KEY AUTOINCREMENT,\n",
    "        category TEXT NOT NULL,\n",
    "        rate REAL NOT NULL\n",
    "    )\n",
    "''')\n",
    "\n",
    "# Data to be inserted into the category_rates table\n",
    "categories = [\n",
    "    ('ADMINISTRATIVE', 0.61),\n",
    "    ('ADVERTISING', 5.31),\n",
    "    ('ARTISAN CONTRACTORS', 2.35),\n",
    "    ('CLUBS/ASSOCIATIONS', 6.53),\n",
    "    ('CONSULTANTS', 1.97),\n",
    "    ('EDUCATION', 9.63),\n",
    "    ('HEALTHCARE', 9.38),\n",
    "    ('OTHER SERVICES', 6.79),\n",
    "    ('PROFESSIONAL SERVICES', 8.22),\n",
    "    ('RESTAURANTS/MOBILE FOOD SERVICES', 2.85),\n",
    "    ('RETAIL', 5.80),\n",
    "    ('TECHNOLOGY', 8.57),\n",
    "    ('WHOLESALE/DISTRIBUTION', 5.61)\n",
    "]\n",
    "\n",
    "# Insert data into the category_rates table\n",
    "cursor.executemany('''\n",
    "    INSERT INTO category_rates (category, rate)\n",
    "    VALUES (?, ?)\n",
    "''', categories)\n",
    "\n",
    "# Commit the changes and close the connection\n",
    "conn.commit()\n",
    "conn.close()\n",
    "\n",
    "print(\"Database and table created successfully, and data inserted.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Required Pydantic Schemas"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "In this section, we define several schemas using Pydantic, a Python library for data validation. \n",
    "\n",
    "These schemas serve as structured models for initial data collection, workflow routing, classification, rationale, and accuracy assessment in our contextual quoting agent. \n",
    "\n",
    "They ensure data consistency and accuracy throughout the quoting process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 201,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel, Field\n",
    "from typing import Optional, List\n",
    "\n",
    "# Core business information model for initial data collection\n",
    "class BusinessInformation(BaseModel):\n",
    "    legal_name: Optional[str] = Field(default=\"Info not yet provided\", description=\"The legal name of the business\")\n",
    "    primary_products_or_services: Optional[str] = Field(default=\"Info not yet provided\", description=\"The primary products or services offered\")\n",
    "    secondary_or_ancillary_operations: Optional[str] = Field(default=\"Info not yet provided\", description=\"Any secondary or ancillary operations\")\n",
    "    industries: Optional[str] = Field(default=\"Info not yet provided\", description=\"In which industries does the business operate?\")\n",
    "    manufacturing_retail_wholesale_or_distribution: Optional[str] = Field(default=\"Info not yet provided\", \n",
    "        description=\"Does the business engage in manufacturing, retail, wholesale, or distribution?\")\n",
    "    projected_revenue: Optional[str] = Field(default=\"Info not yet provided\", description=\"The projected revenue for the upcoming year\")\n",
    "\n",
    "# Models for workflow routing between AI assistants\n",
    "class ToUnderwritingAssistant(BaseModel):\n",
    "    \"\"\"Transfer work to the underwriting assistant\"\"\"\n",
    "\n",
    "    request: str = Field(\n",
    "        description=\"Any information that the underwriting assistant should know before proceeding.\"\n",
    "     )\n",
    "\n",
    "    class Config:\n",
    "        json_schema_extra = {\n",
    "            \"example\": {\n",
    "                \"request\": \"The business is a landscaping company that offers residential and commercial landscaping services, including lawn maintenance, garden design, tree trimming, and irrigation system installation.\"\n",
    "            }\n",
    "        }\n",
    "\n",
    "# Models for classification and code assignment\n",
    "class FinalClassificationsInfo(BaseModel):\n",
    "    \"\"\"Information about the final classifications and code that best describe the business.\"\"\"\n",
    "    \n",
    "    category: str = Field(\n",
    "        description=\"A string representing the category.\"\n",
    "    )\n",
    "    classification_description: str = Field(\n",
    "        description=\"A string representing the category description.\"\n",
    "    )\n",
    "    code: int = Field(\n",
    "        description=\"An integer representing the code assigned to the category.\"\n",
    "    )\n",
    "\n",
    "class ToQuoteAssistant(BaseModel):\n",
    "    \"\"\"Transfer work to the underwriting assistant\"\"\"\n",
    "\n",
    "    request: str = Field(\n",
    "        description=\"Any information that the underwriting assistant should know before proceeding.\"\n",
    "     )\n",
    "    \n",
    "    final_classifications: list[FinalClassificationsInfo] = Field(\n",
    "        description=\"A list of FinalClassificationsInfo objects, each containing a category, category description, and code.\"\n",
    "    )\n",
    "    \n",
    "    class Config:\n",
    "        json_schema_extra = {\n",
    "            \"example\": {\n",
    "                \"request\": \"The business is a landscaping company that offers residential and commercial landscaping services, including lawn maintenance, garden design, tree trimming, and irrigation system installation.\",\n",
    "                \"final_classifications\": [\n",
    "                    {\"category\": \"CONSULTANTS\", \"classification_description\": \"Consultants--Business Services\", \"code\": 74379},\n",
    "                    {\"category\": \"RESTAURANTS/MOBILE FOOD SERVICES\", \"classification_description\": \"Bakeries–Fast Food–Donut Shop.\", \"code\": 10100}\n",
    "                ]\n",
    "            }\n",
    "        }\n",
    "\n",
    "# Workflow control models\n",
    "class CompleteOrEscalate(BaseModel):\n",
    "    \"\"\"A tool to mark the current task as completed and/or to escalate control of the workflow to the main assistant,\n",
    "    who can re-route the dialog based on the user's needs.\"\"\"\n",
    "\n",
    "    cancel: bool = True\n",
    "    reason: str\n",
    "\n",
    "    class Config:\n",
    "        json_schema_extra = {\n",
    "            \"example\": {\n",
    "                \"cancel\": True,\n",
    "                \"reason\": \"Successfully obtained a quote.\",\n",
    "            },\n",
    "            \"example 2\": {\n",
    "                \"cancel\": True,\n",
    "                \"reason\": \"I have fully completed the task.\",\n",
    "            },\n",
    "            \"example 3\": {\n",
    "                \"cancel\": False,\n",
    "                \"reason\": \"I need to search the user's emails or calendar for more information.\",\n",
    "            },\n",
    "        }\n",
    "\n",
    "# Models for classification reasoning and rationale\n",
    "class CategoryRationale(BaseModel):\n",
    "    \"\"\"Represents a category and its rationale or reasoning.\"\"\"\n",
    "    category: str = Field(..., description=\"The name of the category\")\n",
    "    cat_description: str = Field(..., description=\"The description of the category\")\n",
    "    code: int = Field(..., description=\"The unique code for the category\")\n",
    "    rate: float = Field(..., description=\"The rate for the category\")\n",
    "    rationalization: str = Field(..., description=\"A detailed rationalization for the category\")\n",
    "\n",
    "class CategoriesOutput(BaseModel):\n",
    "    \"\"\"Output structure for a list of categories with rationale explanations.\"\"\"\n",
    "    categories: List[CategoryRationale] = Field(..., description=\"List of categories with their rationale explanations\")\n",
    "\n",
    "# Models for classification accuracy assessment\n",
    "class DescriptionGrade(BaseModel):\n",
    "    \"\"\"Schema to define a business description and its corresponding grade.\n",
    "    \n",
    "    Attributes:\n",
    "        description (str): A string representing the business description.\n",
    "        grade (float): A float representing the grade assigned to the business description.\n",
    "    \"\"\"\n",
    "    category: str = Field(\n",
    "        description=\"A string representing the category.\"\n",
    "    )\n",
    "    cat_description: str = Field(\n",
    "        description=\"A string representing the category description.\"\n",
    "    )\n",
    "    code: int = Field(\n",
    "        description=\"An integer representing the code assigned to the category.\"\n",
    "    )\n",
    "    grade: float = Field(\n",
    "        description=\"A float representing the grade assigned to the business description.\"\n",
    "    )\n",
    "\n",
    "class DescriptionGradeSchema(BaseModel):\n",
    "    \"\"\"Schema to hold a list of business descriptions with their corresponding grades.\n",
    "    \n",
    "    Attributes:\n",
    "        description_grades (list[DescriptionGrade]): A list of DescriptionGrade objects, each containing a business description and its corresponding grade.\n",
    "    \"\"\"\n",
    "    description_grades: list[DescriptionGrade] = Field(\n",
    "        description=\"A list of DescriptionGrade objects, each containing a business description and its corresponding grade.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define State"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **State** classes which represent the core data structures for managing and storing various types of information during the workflow. \n",
    "\n",
    "These state definitions ensure that different parts of the quoting process stay synchronized and properly updated:\n",
    "\n",
    "- **State Reducer**: Implements `update_workflow_state`, a function that manages workflow transitions, preventing duplicate entries and handling \"push\" or \"pop\" operations.\n",
    "- **MainState**: Defines the main working state, which includes messages, business information, and workflow state. It maintains critical details about the workflow's progression.\n",
    "- **RAGState**: Represents the state for retrieval-augmented generation (RAG), storing descriptions, documents, rationale, and graded categories to assist with generating and validating responses.\n",
    "- **ExtraState**: Holds additional data, such as tool call identifiers and final classifications, to supplement the workflow process.\n",
    "\n",
    "These states ensure data integrity across all interactions of the contextual quoting agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated, Literal, Optional\n",
    "from typing_extensions import TypedDict\n",
    "from langgraph.graph.message import AnyMessage, add_messages\n",
    "\n",
    "# State Reducer\n",
    "def update_workflow_state(left: list[str], right: Optional[str]) -> list[str]:\n",
    "    \"\"\"Push or pop the workflow state, preventing duplicates.\"\"\"\n",
    "    if right is None or right == []:\n",
    "        return left\n",
    "\n",
    "    # Handle 'pop' case to remove the last element\n",
    "    if right == \"pop\":  \n",
    "        return left[:-1] if left else left  # Safeguard against empty lists\n",
    "    \n",
    "    # Handle cases where 'right' is a list (from the tutorial)\n",
    "    if isinstance(right, list):\n",
    "        # Avoid adding elements already in the list\n",
    "        return left + [item for item in right if item not in left]\n",
    "\n",
    "    # If 'right' is a single string, avoid duplication\n",
    "    if right not in left:\n",
    "        return left + [right]\n",
    "\n",
    "    return left \n",
    "\n",
    "# Main State\n",
    "class MainState(TypedDict):\n",
    "    messages: Annotated[list[AnyMessage], add_messages]\n",
    "    business_information: dict[str, str | float | int]  # Allow values to be str, float, or int\n",
    "    workflow_state: Annotated[\n",
    "        list[\n",
    "            Literal[\n",
    "                \"main_assistant\",\n",
    "                \"underwriting_assistant\",\n",
    "                \"quote_assistant\",\n",
    "            ]\n",
    "        ],\n",
    "        update_workflow_state\n",
    "    ]\n",
    "\n",
    "# RAG State\n",
    "class RAGState(TypedDict):\n",
    "    description: str\n",
    "    documents: list[str]\n",
    "    rationale: str\n",
    "    graded_categories: list[dict[str, float]]\n",
    "\n",
    "# Extra State\n",
    "class ExtraState(TypedDict):\n",
    "    tool_call_id: str\n",
    "    final_classifications: dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define **Tools** that are crucial for enabling specific functions in the contextual quoting agent. Here is a brief explanation:\n",
    "\n",
    "- **Handle Tool Errors**: The function `handle_tool_error()` is designed to manage errors that occur during tool execution. It takes the state, identifies the error, and generates messages indicating what went wrong, allowing users to address and correct the mistakes.\n",
    "\n",
    "- **Create Tool Node with Fallback**: The function `create_tool_node_with_fallback()` creates a tool node using **ToolNode**, and attaches a fallback mechanism using `RunnableLambda` to handle errors gracefully. This ensures that if a tool encounters an error, it defaults to an error handler that provides useful feedback.\n",
    "\n",
    "These definitions provide robustness to the workflow by managing tool errors effectively and ensuring smooth execution through fallback mechanisms.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 203,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.prebuilt import ToolNode\n",
    "from langchain_core.runnables import RunnableLambda\n",
    "from langchain_core.messages import ToolMessage\n",
    "\n",
    "# Handle tool errors\n",
    "def handle_tool_error(state) -> dict:\n",
    "    error = state.get(\"error\")\n",
    "    tool_calls = state[\"messages\"][-1].tool_calls\n",
    "    return {\n",
    "        \"messages\": [\n",
    "            ToolMessage(\n",
    "                content=f\"Error: {repr(error)}\\\\n please fix your mistakes.\",\n",
    "                tool_call_id=tc[\"id\"],\n",
    "            )\n",
    "            for tc in tool_calls\n",
    "        ]\n",
    "    }\n",
    "\n",
    "# Create a tool node with fallback\n",
    "def create_tool_node_with_fallback(tools: list) -> dict:\n",
    "    return ToolNode(tools).with_fallbacks(\n",
    "        [RunnableLambda(handle_tool_error)], exception_key=\"error\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## RAG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **RAG (Retrieval-Augmented Generation)** process, which is used to load and prepare the necessary data for generating responses:\n",
    "\n",
    "- **Load Categories**: We use `TextLoader` to load the category data from a CSV file (`cat_codes_descriptions.csv`). This data contains information about different categories, which is crucial for classification tasks within the quoting agent.\n",
    "\n",
    "- **Split Categories**: We then use `RecursiveCharacterTextSplitter` to split the loaded text into manageable chunks. This splitter divides the text into smaller, non-overlapping chunks based on line breaks, ensuring that each chunk contains a distinct line of data.\n",
    "\n",
    "These steps prepare the data for efficient retrieval, allowing the agent to use relevant pieces of information effectively during the response generation process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 204,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "\n",
    "# Load Categories\n",
    "categories_path = \"data/cat_codes_descriptions.csv\"\n",
    "\n",
    "loader = TextLoader(file_path=categories_path)\n",
    "\n",
    "data_txt = loader.load()\n",
    "\n",
    "# Split Categories\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "        chunk_size=150,  # Smaller chunk size to ensure one line per chunk\n",
    "        chunk_overlap=0,  # No overlap needed for line-by-line splitting\n",
    "        separators=[\"\\n\"],  # Primary separator is line breaks\n",
    "        length_function=len,  # Length function based on character count\n",
    "    )\n",
    "\n",
    "splits = text_splitter.split_documents(data_txt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VectorStore Retriever"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **VectorStore Retriever**, which is used for efficient information retrieval in the quoting agent:\n",
    "\n",
    "- **VectorStore Creation**: We use **Chroma** to create a vector store from the document chunks (`splits`). The vector store acts as a database where the documents are embedded and stored. The **OpenAIEmbeddings** class is used to convert the text chunks into vector representations, making them suitable for similarity-based searches.\n",
    "\n",
    "- **Retriever Setup**: We create a retriever from the vector store using `as_retriever()`, specifying that it should retrieve the top `k=10` most relevant results. This allows the agent to efficiently find related information when handling user queries.\n",
    "\n",
    "This setup enhances the agent's ability to provide accurate, context-aware responses by quickly accessing relevant documents from the vector store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 205,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "# Add to vectorDB\n",
    "vectorstore = Chroma.from_documents(documents=splits,\n",
    "                                    collection_name=\"rag-chroma\",\n",
    "                                    embedding=OpenAIEmbeddings()\n",
    "                                )\n",
    "retriever = vectorstore.as_retriever(search_kwargs={\"k\": 10})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Business Info Gathering Tool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **Business Info Gathering Tool**, which is used to collect and process key business details:\n",
    "\n",
    "- **Business Info Gathering Tool**: The `info_tool` is defined using the `@tool` decorator from **LangChain Core Tools**. This tool takes an instance of `BusinessInformation` as input and processes it to update the state (`MainState`).\n",
    "\n",
    "This tool is essential for capturing and integrating business data into the workflow, ensuring that subsequent steps have access to accurate and complete business information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 206,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.tools import tool\n",
    "\n",
    "# Business Info Gathering Tool\n",
    "@tool\n",
    "def info_tool(business_info: BusinessInformation) -> MainState:\n",
    "    \"\"\"\n",
    "    A tool to process and store business information.\n",
    "    \"\"\"\n",
    "    result = {\n",
    "        \"business_information\": business_info.model_dump()\n",
    "    }\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quote Tool"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **Quote Tool**, which calculates the premium for given categories based on projected revenue:\n",
    "\n",
    "- **Quote Tool Function**: The `fetch_rates` function is decorated with `@tool`, making it available as a tool in the workflow. It takes two arguments:\n",
    "  - `categories`: A list of categories for which the rates are fetched.\n",
    "  - `projected_revenue`: The projected revenue to be used in calculating the premiums.\n",
    "\n",
    "- **Premium Calculation**: \n",
    "  - The function connects to the SQLite database (`categories.db`) to retrieve rates for each provided category.\n",
    "  - It calculates each category's share of the projected revenue and then determines the computed premium based on the rate per $1,000 of revenue.\n",
    "  - The total premium is accumulated across all categories.\n",
    "\n",
    "This tool enables the quoting agent to quickly and accurately provide detailed insurance premium quotes for various business categories, making it a vital component of the system's quoting functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 207,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Quote Tool\n",
    "@tool\n",
    "def fetch_rates(categories: list[str], projected_revenue: float) -> str:\n",
    "    \"\"\"\n",
    "    Fetches the rates for given categories and calculates the computed results based on the projected revenue.\n",
    "\n",
    "    Args:\n",
    "        categories (list[str]): A list of category names for which rates need to be fetched. Example categories include:\n",
    "            - CATEGORY\n",
    "            - ADMINISTRATIVE\n",
    "            - ADVERTISING\n",
    "            - ARTISAN CONTRACTORS\n",
    "            - CLUBS/ASSOCIATIONS\n",
    "            - CONSULTANTS\n",
    "            - EDUCATION\n",
    "            - HEALTHCARE\n",
    "            - OTHER SERVICES\n",
    "            - PROFESSIONAL SERVICES\n",
    "            - RESTAURANTS/MOBILE FOOD SERVICES\n",
    "            - RETAIL\n",
    "            - TECHNOLOGY\n",
    "            - WHOLESALE/DISTRIBUTION\n",
    "        projected_revenue (float): The projected revenue to be distributed among the categories.\n",
    "\n",
    "    Returns:\n",
    "        str: A formatted message containing the quote breakdown\n",
    "    \"\"\"\n",
    "    # Connect to SQLite database\n",
    "    conn = sqlite3.connect('data/categories.db')\n",
    "    cursor = conn.cursor()\n",
    "\n",
    "    # Prepare parts to build the formatted message\n",
    "    message_parts = [\"Here is the requested quote:\\n\", \"Quote Breakdown:\"]\n",
    "\n",
    "    # Initialize total premium\n",
    "    total_premium = 0\n",
    "\n",
    "    # Fetch rate for each classification\n",
    "    for classification in categories:\n",
    "        cursor.execute('''\n",
    "            SELECT rate FROM category_rates WHERE category = ?\n",
    "        ''', (classification,))\n",
    "        result = cursor.fetchone()\n",
    "        if result:\n",
    "            rate = result[0]\n",
    "            portion_revenue = projected_revenue / len(categories)\n",
    "            computed_result = (portion_revenue / 1000) * rate\n",
    "            total_premium += computed_result  # Accumulate total premium\n",
    "\n",
    "            # Add formatted message\n",
    "            message_parts.append(\n",
    "                f\"\\nCategory: {classification}\"\n",
    "                f\"\\n  - Rate: ${rate:.2f} per $1,000\"\n",
    "                f\"\\n  - Portion of Projected Revenue: ${portion_revenue:,.2f}\"\n",
    "                f\"\\n  - Computed Premium: ${computed_result:,.2f}\"\n",
    "            )\n",
    "        else:\n",
    "            message_parts.append(f\"\\nWarning: No rate found for category: {classification}\")\n",
    "\n",
    "    # Add total premium to the message\n",
    "    message_parts.append(f\"\\n\\nTotal Annual Premium: ${total_premium:,.2f}\")\n",
    "\n",
    "    # Close the connection\n",
    "    conn.close()\n",
    "\n",
    "    return \"\\n\".join(message_parts)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Assistants"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define a **Helper Function** that manages interactions with the LLM to streamline assistant behavior:\n",
    "\n",
    "- **Assistant Class**: The `Assistant` class is a helper designed to handle calls to the LLM (Language Learning Model). It is initialized with a `Runnable` object, which represents a configured LLM or an operation that can be executed.\n",
    "\n",
    "This helper function plays a crucial role in maintaining smooth and consistent interactions between the quoting agent and the LLM, handling any cases where the model might fail to produce a satisfactory response initially."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 208,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables import Runnable, RunnableConfig\n",
    "\n",
    "# Helper Assistant Class for handling LLM calls\n",
    "class Assistant:\n",
    "    def __init__(self, runnable: Runnable):\n",
    "        self.runnable = runnable.with_config()\n",
    "        \n",
    "    def __call__(self, state: MainState, config: RunnableConfig):\n",
    "        # Create input with both messages and business_information\n",
    "        input_state = {\n",
    "            \"messages\": state[\"messages\"],\n",
    "            \"business_information\": state.get(\"business_information\", None),\n",
    "        }\n",
    "        \n",
    "        while True:\n",
    "            result = self.runnable.invoke(input_state)\n",
    "            # If the LLM happens to return an empty response, we will re-prompt it\n",
    "            if not result.tool_calls and (\n",
    "                not result.content\n",
    "                or isinstance(result.content, list)\n",
    "                and not result.content[0].get(\"text\")\n",
    "            ):\n",
    "                messages = state[\"messages\"] + [(\"user\", \"Respond with a real output.\")]\n",
    "                state = {**state, \"messages\": messages}\n",
    "                input_state[\"messages\"] = messages\n",
    "            else:\n",
    "                break\n",
    "        return {\"messages\": result}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Main Assistant"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **Main Assistant** of our contextual quoting system. This assistant guides the user through the quoting process, ensuring the gathered data is complete and consistent:\n",
    "\n",
    "- **Main Assistant Prompt**: \n",
    "  - The `main_assistant_prompt` uses `ChatPromptTemplate.from_messages()` to define the conversation flow.\n",
    "  - It defines the assistant's role as a **commercial underwriter assistant**, responsible for guiding the user through the quoting process.\n",
    "  - The assistant's main tasks include:\n",
    "    - Using the `\"info_tool\"` to process and update the business information as the user provides it.\n",
    "    - Evaluating the collected information with underwriters via `\"ToUnderwritingAssistant\"` before proceeding with the quote.\n",
    "    - Providing a summary of the gathered information to the user before moving forward.\n",
    "    - Ensuring only **one tool** is called at a time for better control of actions.\n",
    "\n",
    "- **Main Assistant Tools**: The tools available to the `main_assistant` are defined in `main_assistant_tool`, starting with `[info_tool]`. This ensures the assistant can gather and update business information during the conversation.\n",
    "\n",
    "- **Runnable**:\n",
    "  - `main_assistant_runnable` binds the `main_assistant_prompt` with the LLM (`llm`) using the tools available (`main_assistant_tool + [ToUnderwritingAssistant]`).\n",
    "  - This makes the assistant capable of managing user information collection, calling underwriters, and then proceeding with further steps as needed.\n",
    "\n",
    "This setup allows the **Main Assistant** to effectively manage the quoting workflow, guiding the user while ensuring that the provided information is properly processed, evaluated, and summarized before the next step is taken."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 209,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Main Assistant Prompt\n",
    "main_assistant_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\",\n",
    "        \"\"\"\n",
    "You are the main commercial underwriter assistant.\\n\n",
    "You are responsible for guiding the user through the quoting process and ensuring that the information they provide is accurate and consistent, to later provide a quote. \n",
    "When details and/or information about the business are provided, use \"info_tool\" to process the information and gather or update the required data.\\n\\n\n",
    "Once you have gathered the information, you must first evaluate the prospects information with the underwriters by calling \"ToUnderwritingAssistant\".\\n\n",
    "Before proceeding to quote, you must provide the user with a summary of the information they have provided.\\n\n",
    "**Only call one tool at a time.**\\n\n",
    "Current business information:\\n\n",
    "{business_information}\n",
    "\"\"\",\n",
    "         ),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Main Assistant Tools\n",
    "main_assistant_tool = [info_tool]\n",
    "\n",
    "# Runnable\n",
    "main_assistant_runnable = main_assistant_prompt | llm.bind_tools(main_assistant_tool + [ToUnderwritingAssistant])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Underwriting Assistant"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **Underwriting Assistant**, which is tasked with evaluating the prospect's information and guiding the quoting process after the initial information collection.\n",
    "\n",
    "- **Underwriting Assistant Prompt**:\n",
    "  - The assistant's role is as a **main commercial underwriter**, where its responsibilities include:\n",
    "    - Evaluating the business information provided by the user as the **principal underwriter**.\n",
    "    - Ensuring accuracy in gathered information and confirming relevant business categories.\n",
    "    - Once all the required information has been gathered and validated, the assistant calls `\"ToQuoteAssistant\"` to proceed with the quoting process.\n",
    "\n",
    "The **Underwriting Assistant** plays a critical role in the quoting workflow, acting as a checkpoint to ensure all required information is accurate and sufficient before proceeding to the actual quote generation. This assistant ensures that no important data is missed, maintaining quality and reliability throughout the quoting process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 210,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Underwriting Assistant Prompt\n",
    "underwriting_assistant_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\",\n",
    "        \"\"\"\n",
    "You are the main commercial underwriter assistant.\\n\n",
    "You are responsible for evaluating the prospects information as he principal underwriter.\\n\n",
    "**Only call one tool at a time.**\\n\n",
    "When all the required information and categories have been confirmed, call \"ToQuoteAssistant\" to proceed to quote. \n",
    "Current business information:\\n\n",
    "{business_information}\n",
    "\"\"\",\n",
    "         ),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Runnable\n",
    "underwriting_assistant_runnable = underwriting_assistant_prompt | llm.bind_tools([ToQuoteAssistant])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quote Assistant"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **Quote Assistant**, which is responsible for calculating the insurance premium for a business based on classifications and providing a clear summary to the user.\n",
    "\n",
    "- **Quote Assistant Prompt**:\n",
    "  - The role of the **Quote Assistant** is to act as an expert in underwriting and to generate the **annual premium** based on the business classifications provided.\n",
    "  - Key responsibilities and guidelines for the Quote Assistant include:\n",
    "    1. **Calling `fetch_rates` only once**: It fetches rates for all required categories using the projected revenue.\n",
    "    2. **Finalizing with `CompleteOrEscalate`**: After receiving the rates and calculating the premium, the assistant must call `CompleteOrEscalate` to end the process.\n",
    "    3. **No Recalculation unless instructed**: The assistant should not call `fetch_rates` again unless explicitly asked by the user to do so. \n",
    "    ***(This is to avoid calling the tool multiple times)***\n",
    "  - The assistant also ensures that a clear and understandable **summary** of the quote is provided to the user before concluding.\n",
    "\n",
    "- **Quote Assistant Tools**: \n",
    "  - The tools available to the quote assistant are defined in `quote_assistant_tools`, which includes `[fetch_rates]`. This enables the assistant to fetch rates for given categories and calculate the required premium.\n",
    "\n",
    "The **Quote Assistant** plays a vital role in completing the underwriting and quoting process. By fetching rates, calculating premiums, and providing a summary to the user, it ensures that the entire workflow is concluded seamlessly with all necessary information clearly communicated to the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 211,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Quote Assistant Prompt\n",
    "quote_prompt = ChatPromptTemplate.from_messages([\n",
    "    (\"system\",\n",
    "    \"\"\"\n",
    "You are an expert commercial underwriting assistant tasked with providing a quote for the business based on the classifications.\n",
    "\n",
    "You will use the provided classifications and their corresponding rates to calculate the annual premium.\n",
    "\n",
    "IMPORTANT: \n",
    "1. Call fetch_rates ONLY ONCE with all required categories and the projected revenue.\n",
    "2. After receiving the rates and calculating the premium, call CompleteOrEscalate to end the process.\n",
    "3. Do not call fetch_rates again unless explicitly asked to recalculate.\n",
    "\n",
    "Business Description:\n",
    "{business_information}\n",
    "\n",
    "When you have the quote results, provide a clear summary to the user and then call CompleteOrEscalate with cancel=true to complete the process.\n",
    "\"\"\",\n",
    "         ),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Quote Assistant Tools\n",
    "quote_assistant_tools = [fetch_rates]\n",
    "\n",
    "# Runnable\n",
    "quote_assistant_runnable = quote_prompt | llm.bind_tools(quote_assistant_tools + [CompleteOrEscalate])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agents "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rewriting Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **Rewriting Agent**, which is responsible for extracting and summarizing key operational details from business descriptions to support the insurance underwriting process.\n",
    "\n",
    "**Different from Assistants, Agents are used to process information and return structured outputs.**\n",
    "**They do not interact with the user.**\n",
    "\n",
    "- **Retriever Assistant Prompt**:\n",
    "  - The assistant's role is to analyze the given **business description** and extract important operational details relevant to insurance underwriting.\n",
    "  - Key tasks include:\n",
    "    1. **Analyze** the provided business description.\n",
    "    2. **Identify and Summarize** key aspects like workflows, operational tasks, and products/services offered.\n",
    "    3. **Highlight** specific activities that are critical for matching the business to the appropriate **code classification**.\n",
    "  - The summary should be **concise** and **structured** with bullet points for easy readability, as shown in the provided example.\n",
    "\n",
    "The **Rewriting Agent** helps ensure that business descriptions are processed into meaningful insights that can assist in determining appropriate insurance classifications. This agent makes the underwriting process more streamlined by providing clear and actionable information derived from user-provided data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 212,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "# Retriever Assistant Prompt\n",
    "retriever_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"\"\"\n",
    "You are an expert commercial underwriting assistant specialized in analyzing business descriptions to extract key operational details relevant to the insurance underwriting process.\n",
    "\n",
    "**Your task is to:**\n",
    "\n",
    "1. **Analyze** the provided business description.\n",
    "2. **Identify and summarize** the key aspects of the business operations, focusing on workflows, operational tasks, products/services offered, and any other relevant details.\n",
    "3. **Highlight** any specific activities that are crucial for matching the business to the appropriate insurance code classification.\n",
    "\n",
    "**Please provide the summary in a clear, concise, and structured format, using bullet points for each key aspect.**\n",
    "\n",
    "---\n",
    "\n",
    "**Business Information:**\n",
    "\n",
    "{business_information}\n",
    "\n",
    "---\n",
    "\n",
    "**Example:**\n",
    "\n",
    "**Business Description:**\n",
    "\n",
    "\"**ABC Landscaping** offers residential and commercial landscaping services, including lawn maintenance, garden design, tree trimming, and irrigation system installation. They also provide seasonal services like snow removal. The company operates a fleet of service vehicles and heavy machinery such as lawn mowers, trimmers, and excavators.\"\n",
    "\n",
    "**Summary:**\n",
    "\n",
    "- Provides residential and commercial landscaping services.\n",
    "- Services include lawn maintenance, garden design, and tree trimming.\n",
    "- Offers installation of irrigation systems.\n",
    "- Provides seasonal snow removal services.\n",
    "- Operates a fleet of service vehicles and heavy machinery (lawn mowers, trimmers, excavators).\n",
    "\n",
    "\"\"\",\n",
    "        )\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Runnable\n",
    "retriever_runnable = retriever_prompt | llm | StrOutputParser()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reasoning Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **Reasoning Agent**, which is tasked with determining the relevance of different insurance categories for a business based on its description.\n",
    "\n",
    "- **Reasoning Categories Assistant Prompt**:\n",
    "  - The role of the **Reasoning Agent** is to evaluate the **business description** and determine which categories apply to the nature of the operations.\n",
    "  - The agent uses **logical reasoning** to assess which categories from the given list are relevant.\n",
    "  - The assistant receives both the **business description** and a list of possible **categories** and provides a reasoned analysis of which ones are applicable.\n",
    "  - The response includes not just the category, but also a **rationalization**—a brief explanation of why a particular category was chosen based on the business operations. The provided example demonstrates how to present the selected categories with reasoning.\n",
    "\n",
    "The **Reasoning Agent** is critical for ensuring that the business operations are properly categorized with a logical and justifiable approach. By providing reasoned categorization, it helps ensure accuracy and compliance in matching businesses to the right insurance classification, which is key for accurate underwriting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 213,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reasoning Categories Assistant Prompt\n",
    "reasoning_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\",\n",
    "        \"\"\"\n",
    "You are an expert commercial underwriting assistant tasked with providing reasoning for the categories of the nature of operations of a business.\\n\n",
    "You will be given a description of the business and a list of categories that may apply to the business.\\n\n",
    "Use logical reasoning to determine which categories are relevant to the business based on the provided description.\\n\n",
    "Business Description:\\n\n",
    "{description}\\n\n",
    "Categories:\\n\n",
    "{documents}\\n\n",
    "\n",
    "---\n",
    "\n",
    "**Example:**\n",
    "\n",
    "\n",
    "{{\n",
    "    \"categories\": [\n",
    "      {{\n",
    "        \"category\": \"CONSULTANTS\",\n",
    "        \"description\": \"Consultants--Business Services\"\n",
    "        \"code\": \"74379\",\n",
    "        \"rate\": \"1.97\",\n",
    "        \"rationalization\": \"GABAS, LLC provides consulting services related to AI and technology, which falls under the category of business services. This includes advising businesses on AI implementation and strategy.\"\n",
    "      }},\n",
    "      {{\n",
    "        \"category\": \"RESTAURANTS/MOBILE FOOD SERVICES\",\n",
    "        \"description\": \"Bakeries–Fast Food–Donut Shop.\"\n",
    "        \"code\": \"10100\",\n",
    "        \"rate\": \"2.85\",\n",
    "        \"rationalization\": \"A bakery that offers fast food and donut services is considered part of the food service industry.\"\n",
    "      }}\n",
    "    ]\n",
    "}}\n",
    "\n",
    "\"\"\",\n",
    "         )\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Runnable\n",
    "reasoning_runnable = reasoning_prompt | llm.with_structured_output(CategoriesOutput)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classification Grading Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we define the **Classification Grading Agent**, which evaluates the accuracy of category assignments for a business based on its description.\n",
    "\n",
    "- **Classification Grading Assistant Prompt**:\n",
    "  - The **Classification Grading Agent** is responsible for assessing the match between the **business description** and the **assigned categories**.\n",
    "  - The agent uses logical reasoning to **grade the accuracy** of the assigned categories to ensure they are the most relevant and appropriate for the business operations.\n",
    "  - The prompt includes:\n",
    "    - A **business description**.\n",
    "    - The **category descriptions and rationalizations** (`{cat_and_rat}`) to provide context for the grading process.\n",
    "  - The goal is to ensure that each classification is graded based on how well it aligns with the details of the business.\n",
    "\n",
    "- **Runnable**:\n",
    "  - `classification_grading_runnable` binds the `classification_grading_prompt` with the LLM (`llm`) and uses `.with_structured_output(DescriptionGradeSchema)` to ensure that the output is returned in a structured format.\n",
    "  - The **DescriptionGradeSchema** includes details about the category, its description, and the assigned grade, helping evaluate the fit between the business operations and the assigned categories.\n",
    "\n",
    "The **Classification Grading Agent** is an important step in the underwriting process, helping validate the quality of the category assignments. It ensures that each classification matches the actual operations of the business, which is key for accurate risk assessment and premium calculation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 214,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Classification Grading Assistant Prompt\n",
    "classification_grading_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\",\n",
    "        \"\"\"\n",
    "You are an expert commercial underwriting assistant tasked with grading the accuracy between a given business description and the classifications assigned.\\n\n",
    "Apply logical reasoning, grade the accuracy between the business description and the categories assigned.\\n\n",
    "The objective is to match the business description with the most relevant categories.\\n\n",
    "Business Description:\\n\n",
    "{description}\\n\n",
    "Category Description and Rationalization:\\n\n",
    "{cat_and_rat}\n",
    "\"\"\",\n",
    "         )\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Runnable\n",
    "classification_grading_runnable = classification_grading_prompt | llm.with_structured_output(DescriptionGradeSchema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define Nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Retriever Node"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define a **Retriever Node** named `reasoning`, which is responsible for providing logical reasoning to identify relevant categories for a business based on its description and supporting documents.\n",
    "\n",
    "- **Reasoning Node Function**:\n",
    "  - The function `reasoning(rag_state: RAGState) -> RAGState` takes an `RAGState` as input, which includes the business **description** and the **retrieved documents** (likely possible categories).\n",
    "  - **Reasoning Process**:\n",
    "    - The function uses the **Reasoning Agent** (`reasoning_runnable`) to generate a rationale based on the provided description and documents.\n",
    "    - The `invoke()` method is called with a dictionary containing the **description** and **documents**, which the reasoning agent processes to provide the rationale.\n",
    "  - The function then returns an updated `RAGState`, which now includes the reasoned **categories** of the business, stored in `\"rationale\"`.\n",
    "\n",
    "The **Retriever Node** adds depth to the workflow by analyzing the business description alongside related documents to provide detailed reasoning behind category choices. This helps ensure that all classifications assigned to the business are thoroughly validated, contributing to a more accurate and justified underwriting process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 215,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retrieve Node\n",
    "def retrieve(state: MainState) -> RAGState:\n",
    "    \"\"\"\n",
    "    The function will first generate a detailed description of the business and then retrieve relevant documents.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): New key added to state, documents, that contains retrieved documents\n",
    "    \"\"\"\n",
    "    print(\"---Rewrite Description---\")\n",
    "    nature_of_ops = state[\"business_information\"]\n",
    "    description = retriever_runnable.invoke({\"business_information\": nature_of_ops})\n",
    "    print(f\"Description:\\n {description}\")\n",
    "\n",
    "    print(\"---RETRIEVE---\")\n",
    "    # Retrieval\n",
    "    documents = retriever.invoke(description)\n",
    "    print(f\"Documents:\\n {documents}\")\n",
    "    return {\"description\": description, \"documents\": documents}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reasoning Node"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the **Reasoning Node**, which is responsible for analyzing the business description and providing a reasoned rationale for the assigned categories.\n",
    "\n",
    "- **Reasoning Node Function**:\n",
    "  - The function `reasoning(rag_state: RAGState) -> RAGState` is designed to process the **RAGState** containing:\n",
    "    - The **description** of the business.\n",
    "    - The **retrieved documents**, which typically represent possible categories for classification.\n",
    "  - **Reasoning Process**:\n",
    "    - The function calls the **Reasoning Agent** (`reasoning_runnable`) using the `invoke()` method, passing in the `description` and `documents` to generate the **rationale**.\n",
    "    - This rationale explains why certain categories are relevant to the business, providing a more thorough understanding of the classification logic.\n",
    "  - **Return Value**: \n",
    "    - The function returns an updated `RAGState`, including the reasoned **rationale** for the selected categories (`\"rationale\": rationale`).\n",
    "\n",
    "The **Reasoning Node** is crucial for adding transparency and justification to the underwriting process. By providing a detailed explanation of why certain categories are assigned, it ensures that the classification is not only accurate but also understandable, making it easier for users or underwriters to validate and trust the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 216,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reasoning Node\n",
    "def reasoning(rag_state: RAGState) -> RAGState:\n",
    "    \"\"\"\n",
    "    The function provides reasoning for the business categories based on the provided description and retrieved documents.\n",
    "\n",
    "    Args:\n",
    "        rag_state (RAGState): The current state containing the description and retrieved documents.\n",
    "\n",
    "    Returns:\n",
    "        RAGState: Updated state with reasoned categories of the business.\n",
    "    \"\"\"\n",
    "    print(\"---REASONING---\")\n",
    "    description = rag_state[\"description\"]\n",
    "    documents = rag_state[\"documents\"]\n",
    "\n",
    "    # Reasoning\n",
    "    rationale = reasoning_runnable.invoke({\"description\": description, \"documents\": documents})\n",
    "    print(f\"Rationale:\\n {rationale}\")\n",
    "    return {\"rationale\": rationale}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Classification Grading Node"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the **Classification Grading Node**, which evaluates how well each business category matches the provided business description.\n",
    "\n",
    "- **Classification Grading Node Function**:\n",
    "  - The `classification_grading(rag_state: RAGState) -> ExtraState` function is designed to **grade** each of the business categories provided in the `RAGState`.\n",
    "  - **Core Steps**:\n",
    "    1. **Extract Core Data**:\n",
    "       - Extracts the **description** of the business, **rationale**, and the **categories** from the provided `rag_state`.\n",
    "       - Prints the number of categories being processed for grading.\n",
    "    2. **Formatting for Grading**:\n",
    "       - Formats each category's information (category name, description, code, and rationalization) into a readable string called `cat_and_rat`.\n",
    "       - This formatted string is used as input to the **Classification Grading Agent** (`classification_grading_runnable`) to evaluate the categories.\n",
    "    3. **Grading Categories**:\n",
    "       - Calls `classification_grading_runnable.invoke()` with the business **description** and the formatted `cat_and_rat` to grade the accuracy of the assigned categories.\n",
    "       - The output (`graded_categories`) includes a graded evaluation of each category.\n",
    "  - **Return Value**:\n",
    "    - The function returns an `ExtraState` that includes the **graded categories** (`\"graded_categories\": graded_categories.description_grades`).\n",
    "\n",
    "The **Classification Grading Node** is essential for ensuring that the assigned categories accurately reflect the business operations. It helps validate and refine the classification process, providing a clear and justifiable grading that can support underwriting decisions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 217,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Classification Grading Node\n",
    "def classification_grading(rag_state: RAGState) -> ExtraState:\n",
    "    \"\"\"\n",
    "    Grade each category based on how well it matches the business description.\n",
    "    \"\"\"\n",
    "    print(\"\\n=== CLASSIFICATION GRADING START ===\")\n",
    "    \n",
    "    # Extract core data\n",
    "    description = rag_state[\"description\"]\n",
    "    rationale = rag_state[\"rationale\"]\n",
    "    categories = rationale.categories\n",
    "    \n",
    "    print(f\"Processing {len(categories)} categories for grading...\")\n",
    "    \n",
    "    try:\n",
    "        # Format categories into a single string\n",
    "        cat_and_rat = \"\\n\\n\".join([\n",
    "            f\"Category: {cat.category}\\n\"\n",
    "            f\"Description: {cat.cat_description}\\n\"\n",
    "            f\"Code: {cat.code}\\n\"\n",
    "            f\"Rationale: {cat.rationalization}\\n\"\n",
    "            for cat in categories\n",
    "        ])\n",
    "        \n",
    "        graded_categories = classification_grading_runnable.invoke({\n",
    "            \"description\": description,\n",
    "            \"cat_and_rat\": cat_and_rat\n",
    "        })\n",
    "        \n",
    "        print(f\"\\nCompleted grading {len(graded_categories.description_grades)} categories\")\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"Error during classification grading: {str(e)}\")\n",
    "        raise\n",
    "    \n",
    "    print(\"=== CLASSIFICATION GRADING END ===\\n\")\n",
    "    return {\"graded_categories\": graded_categories.description_grades}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Routing Nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We define the **Routing Nodes**, which are responsible for determining the flow of conversation between different assistants within the system. These routing nodes help manage transitions across various stages of the workflow to ensure a seamless process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. **Route Main Assistant (`route_main_assistant`)**:\n",
    "   - Determines the next step for the **main assistant**.\n",
    "   - Depending on the last tool call in the state, the route directs either to continue with tools (`\"main_assistant_tools\"`), retrieve more information (`\"retrieve\"`), or end the session (`\"__end__\"`).\n",
    "   - Routes to `\"pass_tool_call_id\"` if the `ToUnderwritingAssistant` is called, otherwise `\"main_assistant_tools\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 218,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.prebuilt import tools_condition\n",
    "from typing import Literal, Callable\n",
    "\n",
    "# Function to determine the next route for the main assistant\n",
    "def route_main_assistant(\n",
    "    state: MainState,\n",
    ") -> Literal[\n",
    "    \"main_assistant_tools\",\n",
    "    \"retrieve\",\n",
    "    \"__end__\",\n",
    "]:\n",
    "    route = tools_condition(state)\n",
    "    if route == END:\n",
    "        return END\n",
    "    tool_calls = state[\"messages\"][-1].tool_calls\n",
    "    if tool_calls:\n",
    "        if tool_calls[0][\"name\"] == ToUnderwritingAssistant.__name__:\n",
    "            return \"pass_tool_call_id\"\n",
    "        else:\n",
    "            return \"main_assistant_tools\"\n",
    "    raise ValueError(\"Invalid route\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. **Route Underwriting Assistant (`route_underwriting_assistant`)**:\n",
    "   - Manages the flow for the **underwriting assistant**.\n",
    "   - If the `ToQuoteAssistant` tool is called, it routes to `\"update_workflow_state\"`, allowing progression to the next step in the quoting process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 219,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to determine the next route for the underwriting assistant\n",
    "def route_underwriting_assistant(\n",
    "    state: MainState,\n",
    ") -> Literal[\n",
    "\n",
    "    \"update_workflow_state\",\n",
    "    \"__end__\",\n",
    "]:\n",
    "    route = tools_condition(state)\n",
    "    if route == END:\n",
    "        return END\n",
    "    tool_calls = state[\"messages\"][-1].tool_calls\n",
    "\n",
    "    if tool_calls:\n",
    "        if tool_calls[0][\"name\"] == ToQuoteAssistant.__name__:\n",
    "            return \"update_workflow_state\"\n",
    "\n",
    "    raise ValueError(\"Invalid route\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. **Route Quote Assistant (`route_quote_assistant`)**:\n",
    "   - Determines the routing for the **quote assistant**.\n",
    "   - If the conversation is complete and `\"CompleteOrEscalate\"` is called, it reroutes (`\"reroute\"`).\n",
    "   - Otherwise, if there are tools to be used, it directs the flow to `\"quote_assistant_tools\"`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 220,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to determine the next route for the quote assistant\n",
    "def route_quote_assistant(\n",
    "    state: MainState,\n",
    ") -> Literal[\n",
    "    \"quote_assistant_tools\",\n",
    "    \"reroute\",\n",
    "    \"__end__\",\n",
    "]:\n",
    "    route = tools_condition(state)\n",
    "    if route == END:\n",
    "        return END\n",
    "\n",
    "    tool_calls = state[\"messages\"][-1].tool_calls\n",
    "    \n",
    "    did_cancel = any(tc[\"name\"] == CompleteOrEscalate.__name__ for tc in tool_calls)\n",
    "    if did_cancel:\n",
    "        return \"reroute\"\n",
    "    \n",
    "    # Get tool names using the name property for StructuredTools\n",
    "    tool_names = [t.name for t in quote_assistant_tools]\n",
    "    \n",
    "    if all(tc[\"name\"] in tool_names for tc in tool_calls):\n",
    "        return \"quote_assistant_tools\"\n",
    "        \n",
    "    raise ValueError(f\"Invalid route. Tool call {[tc['name'] for tc in tool_calls]} not found in available tools {tool_names}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. **Create Entry Node for Quote Conversation (`create_entry_node_quote`)**:\n",
    "   - Creates an **entry node** for a new quote.\n",
    "   - Updates the state with the `\"final_classifications\"` and provides a message to initiate the quoting process with the selected assistant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 221,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to create an entry node for quote conversation\n",
    "def create_entry_node_quote(assistant_name: str, new_workflow_state: str) -> Callable:\n",
    "    def entry_node(state: MainState):\n",
    "        tool_call_id = state[\"messages\"][-1].tool_calls[0][\"id\"]\n",
    "        final_classifications = state[\"messages\"][-1].tool_calls[0][\"args\"][\"final_classifications\"]\n",
    "        business_information = state.get(\"business_information\", {})  # Changed from [] to {}\n",
    "        \n",
    "        # Create a new dictionary combining both pieces of information\n",
    "        updated_business_info = {\n",
    "            **business_information,\n",
    "            \"final_classifications\": final_classifications\n",
    "        }\n",
    "        return {\n",
    "            \"messages\": [\n",
    "                ToolMessage(\n",
    "                    content=f\"The assistant is now {assistant_name}. You will proceed to provide a quote using the final classifications related to the business in evaluation.\"\n",
    "                    f\" Final Classifications/Categories: {final_classifications}.\"\n",
    "                    \" Do not mention who you are - just act as the proxy for the assistant.\",\n",
    "                    tool_call_id=tool_call_id,\n",
    "                )\n",
    "            ],\n",
    "            \"business_information\": updated_business_info,\n",
    "            \"workflow_state\": new_workflow_state\n",
    "        }\n",
    "\n",
    "    return entry_node\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. **Route to Workflow (`route_to_workflow`)**:\n",
    "   - Determines which assistant to route to based on the current `\"workflow_state\"`.\n",
    "   - Routes back to the **main assistant** if no specific workflow state is set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 222,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to return to the workflow\n",
    "def route_to_workflow(\n",
    "    state: MainState,\n",
    ") -> Literal[\n",
    "    \"main_assistant\",\n",
    "    \"underwriting_assistant\",\n",
    "    \"quote_assistant\",\n",
    "]:\n",
    "    \"\"\"If we are in a workflow state, route directly to the appropriate assistant.\"\"\"\n",
    "    workflow_state = state.get(\"workflow_state\")\n",
    "    if not workflow_state:\n",
    "        return \"main_assistant\"\n",
    "    return workflow_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6. **Reroute Workflow to Main Assistant (`reroute`)**:\n",
    "   - Reroutes the workflow to the **main assistant**, resuming dialog and popping the current workflow state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 223,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to reroute the workflow to the main assistant\n",
    "def reroute(state: MainState) -> dict:\n",
    "    \"\"\"Reroute the workflow to the main assistant.\"\"\"\n",
    "    messages = []\n",
    "    if state[\"messages\"][-1].tool_calls:\n",
    "        messages.append(\n",
    "            ToolMessage(\n",
    "                content=\"Resuming dialog with the main assistant.\",\n",
    "                tool_call_id=state[\"messages\"][-1].tool_calls[0][\"id\"],\n",
    "            )\n",
    "        )\n",
    "    return {\n",
    "        \"workflow_state\": \"pop\",\n",
    "        \"messages\": messages,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7. **Update Workflow State (`update_workflow_state`)**:\n",
    "   - Updates the workflow state by popping the last entry to ensure the correct assistant resumes the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 224,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to update the workflow state\n",
    "def update_workflow_state(state: MainState) -> MainState:\n",
    "    \"\"\"Reroute the workflow to the main assistant.\"\"\"\n",
    "    return {\n",
    "        \"workflow_state\": \"pop\",\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "8. **Pass Tool Call ID to ExtraState (`pass_tool_call_id`)**:\n",
    "   - Transfers the **tool call ID** from the state to `ExtraState` for later use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 225,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to pass the tool call ID to ExtraState\n",
    "def pass_tool_call_id(state: MainState) -> ExtraState:\n",
    "    tool_call_id = state[\"messages\"][-1].tool_calls[0][\"id\"]\n",
    "    return {\"tool_call_id\": tool_call_id}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "9. **Pass Final Classifications (`pass_final_classifications`)**:\n",
    "   - Transfers the **final classifications** from `RAGState` to `ExtraState`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 226,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to pass the final classifications to ExtraState\n",
    "def pass_final_classifications(rag_state: RAGState) -> ExtraState:\n",
    "    final_classifications = rag_state[\"graded_categories\"]\n",
    "    return {\"final_classifications\": final_classifications}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "10. **Create Tool Message (`create_tool_message`)**:\n",
    "    - Generates a **ToolMessage** for the main underwriting assistant.\n",
    "    - It includes information about the identified categories and instructs the assistant to confirm their relevance through follow-up questions without sharing the grades with the user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 227,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to create a tool message and pass to MainState\n",
    "def create_tool_message(extra_state: ExtraState) -> MainState:\n",
    "    tool_call_id = extra_state[\"tool_call_id\"]\n",
    "    final_classifications = extra_state[\"final_classifications\"]\n",
    "    return {\n",
    "            \"messages\": [\n",
    "                ToolMessage(\n",
    "                    content=\"\"f\"Assume the role of the Main Underwriting Assistant and reflect on the previous conversation between the host assistant and the user.\"\n",
    "                    f\" Based on the initial underwriting analysis, the following categories have been identified as most relevant to the business:\"\n",
    "                    f\"{final_classifications}\"\n",
    "                    \"Grades have been assigned based on the accuracy of the descriptions and the rationale provided for each category. KEEP THE GRADES PRIVATE AND DO NOT SHARE THEM WITH THE USER.\"\n",
    "                    \"To confirm their relevance to the business, please ask the user follow-up questions to finalize the validity of these categories.\"\"\",\n",
    "                    tool_call_id=tool_call_id,\n",
    "                )\n",
    "            ],\n",
    "            \"workflow_state\": \"underwriting_assistant\"\n",
    "        }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "11. **Update State (`update_state`)**:\n",
    "    - Updates the main state with new business information retrieved from tool calls.\n",
    "    - Handles potential JSON parsing issues to ensure that the state is updated consistently with accurate information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 228,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to update the state with the new business information from tool calls\n",
    "def update_state(state: MainState) -> MainState:\n",
    "    \"\"\"\n",
    "    The function updates the state with the new business information from tool calls.\n",
    "    \"\"\"\n",
    "    messages = state.get(\"messages\", [])\n",
    "    \n",
    "    # Find the most recent tool message\n",
    "    tool_messages = [msg for msg in messages if msg.type == \"tool\"]\n",
    "    if tool_messages:\n",
    "        latest_tool_message = tool_messages[-1]\n",
    "        # Parse the JSON content from the tool message\n",
    "        import json\n",
    "        try:\n",
    "            tool_content = json.loads(latest_tool_message.content)\n",
    "            business_info = tool_content.get(\"business_information\", {})\n",
    "        except json.JSONDecodeError:\n",
    "            business_info = state.get(\"business_information\", {})\n",
    "    else:\n",
    "        business_info = state.get(\"business_information\", {})\n",
    "\n",
    "    return {\n",
    "        \"messages\": messages,\n",
    "        \"business_information\": business_info\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The **Routing Nodes** play a key role in managing the workflow's progression through different assistants and tools. They handle conditions such as tool usage, completion of processes, and transitions between assistants, ensuring each stage of the quoting and underwriting process flows smoothly and logically. These nodes make the system more modular and efficient by directing the appropriate assistant to continue based on the context and current state of the workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define and Configure Graph"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Graph defines how different parts of the assistant workflow interconnect. Each node represents a specific step or role, and edges dictate the possible transitions between these steps. By adding conditions and sequential edges, the graph provides flexibility in how each assistant and tool works together, ensuring that the entire quoting process is both adaptable and logically structured.\n",
    "\n",
    "This comprehensive setup allows for smooth transitions between data collection, underwriting evaluation, and premium quoting, while also handling possible errors or rerouting needs effectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 229,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import START, StateGraph, END\n",
    "\n",
    "# Graph State\n",
    "graph = StateGraph(MainState)\n",
    "\n",
    "## Add Nodes\n",
    "# Main Assistant Node\n",
    "graph.add_node(\"main_assistant\", Assistant(main_assistant_runnable))\n",
    "# Main Assistant Tools Node\n",
    "graph.add_node(\"main_assistant_tools\", create_tool_node_with_fallback(main_assistant_tool))\n",
    "# Underwriting Assistant Node\n",
    "graph.add_node(\"underwriting_assistant\", Assistant(underwriting_assistant_runnable))\n",
    "# Quote Assistant Node\n",
    "graph.add_node(\"quote_assistant\", Assistant(quote_assistant_runnable))\n",
    "# Quote Assistant Tools Node\n",
    "graph.add_node(\"quote_assistant_tools\", create_tool_node_with_fallback(quote_assistant_tools))\n",
    "# Entry Quote Assistant Node\n",
    "graph.add_node(\n",
    "    \"entry_quote_assistant\",\n",
    "    create_entry_node_quote(\"Main Quote Assistant\", \"quote_assistant\")\n",
    ")\n",
    "# Retrieve Node\n",
    "graph.add_node(\"retrieve\", retrieve)\n",
    "# Reasoning Node\n",
    "graph.add_node(\"reasoning\", reasoning)\n",
    "# Classification Grading Node\n",
    "graph.add_node(\"classification_grading\", classification_grading)\n",
    "# Update State Node\n",
    "graph.add_node(\"update_state\", update_state)\n",
    "# Reroute Node\n",
    "graph.add_node(\"reroute\", reroute)\n",
    "# Pass Tool Call ID Node\n",
    "graph.add_node(\"pass_tool_call_id\", pass_tool_call_id)\n",
    "# Pass Final Classifications Node\n",
    "graph.add_node(\"pass_final_classifications\", pass_final_classifications)\n",
    "# Create Tool Message Node\n",
    "graph.add_node(\"create_tool_message\", create_tool_message)\n",
    "# Update Workflow State Node\n",
    "graph.add_node(\"update_workflow_state\", update_workflow_state)\n",
    "\n",
    "## Add Edges\n",
    "# Set Conditional Entry Point\n",
    "graph.set_conditional_entry_point(\n",
    "    route_to_workflow,\n",
    "    {\n",
    "        \"main_assistant\": \"main_assistant\",\n",
    "        \"underwriting_assistant\": \"underwriting_assistant\",\n",
    "        \"quote_assistant\": \"quote_assistant\",\n",
    "    }\n",
    ")\n",
    "\n",
    "# Conditional Edges from main assistant\n",
    "graph.add_conditional_edges(\n",
    "    \"main_assistant\",\n",
    "    route_main_assistant,\n",
    "    [\n",
    "        \"main_assistant_tools\",\n",
    "        \"pass_tool_call_id\",\n",
    "        END,\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Conditional Edges from underwriting assistant\n",
    "graph.add_conditional_edges(\n",
    "    \"underwriting_assistant\",\n",
    "    route_underwriting_assistant,\n",
    "    [\n",
    "        \"reroute\",\n",
    "        \"update_workflow_state\",\n",
    "        END,\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Conditional Edges from quote assistant\n",
    "graph.add_conditional_edges(\n",
    "    \"quote_assistant\",\n",
    "    route_quote_assistant,\n",
    "    [\n",
    "        \"quote_assistant_tools\",\n",
    "        \"reroute\",\n",
    "        END,\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Edges from update workflow state to entry quote assistant\n",
    "graph.add_edge(\"update_workflow_state\", \"entry_quote_assistant\")\n",
    "# Edges from pass tool call id to retrieve\n",
    "graph.add_edge(\"pass_tool_call_id\", \"retrieve\")\n",
    "# Edges from retrieve to reasoning\n",
    "graph.add_edge(\"retrieve\", \"reasoning\")\n",
    "# Edges from reasoning to classification grading\n",
    "graph.add_edge(\"reasoning\", \"classification_grading\")\n",
    "# Edges from classification grading to pass final classifications\n",
    "graph.add_edge(\"classification_grading\", \"pass_final_classifications\")\n",
    "# Edges from pass final classifications to create tool message  \n",
    "graph.add_edge(\"pass_final_classifications\", \"create_tool_message\")\n",
    "# Edges from create tool message to underwriting assistant\n",
    "graph.add_edge(\"create_tool_message\", \"underwriting_assistant\")\n",
    "# Edges from update state to main assistant\n",
    "graph.add_edge(\"main_assistant_tools\", \"update_state\")\n",
    "# Edges from update state to main assistant\n",
    "graph.add_edge(\"update_state\", \"main_assistant\")\n",
    "# Edges from reroute to main assistant\n",
    "graph.add_edge(\"reroute\", \"main_assistant\")\n",
    "# Edges from entry quote assistant to quote assistant\n",
    "graph.add_edge(\"entry_quote_assistant\", \"quote_assistant\")\n",
    "# Edges from quote assistant tools to quote assistant\n",
    "graph.add_edge(\"quote_assistant_tools\", \"quote_assistant\")\n",
    "\n",
    "# Memory Saver\n",
    "memory = MemorySaver()\n",
    "\n",
    "# Compile Graph\n",
    "final_graph = graph.compile(\n",
    "    checkpointer=memory,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graph Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Final Graph](../images/contextual_quoting_graph.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Quote Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 230,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "GABAS, LLC. We offer tailor made AI solutions. We also provide consulting services. We operate in the tech industry. As software providers, I believe we fall under the distribution type of business. We expect $4 million in revenues.\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  info_tool (77aec3c3-9632-492d-8314-ff7f69f2b362)\n",
      " Call ID: 77aec3c3-9632-492d-8314-ff7f69f2b362\n",
      "  Args:\n",
      "    business_info: {'company_name': 'GABAS, LLC', 'description': ['tailor made AI solutions', 'consulting services'], 'industry': 'tech', 'business_type': 'distribution', 'revenue': 4000000}\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: info_tool\n",
      "\n",
      "Error: 1 validation error for info_tool\n",
      "business_info\n",
      "  Input should be a valid dictionary or instance of BusinessInformation [type=model_type, input_value=\"{'company_name': 'GABAS,...n', 'revenue': 4000000}\", input_type=str]\n",
      "    For further information visit https://errors.pydantic.dev/2.9/v/model_type\n",
      " Please fix your mistakes.\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  info_tool (6fca9f47-e7f1-423a-a44c-d5bcbd6b0aa4)\n",
      " Call ID: 6fca9f47-e7f1-423a-a44c-d5bcbd6b0aa4\n",
      "  Args:\n",
      "    business_info: {'company_name': 'GABAS, LLC', 'description': ['tailor made AI solutions', 'consulting services'], 'industry': 'tech', 'business_type': 'software_provider', 'revenue': 4000000}\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: info_tool\n",
      "\n",
      "Error: 1 validation error for info_tool\n",
      "business_info\n",
      "  Input should be a valid dictionary or instance of BusinessInformation [type=model_type, input_value=\"{'company_name': 'GABAS,...r', 'revenue': 4000000}\", input_type=str]\n",
      "    For further information visit https://errors.pydantic.dev/2.9/v/model_type\n",
      " Please fix your mistakes.\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  info_tool (ed7edf8b-e21c-4dcf-bc53-b9b7f1ce444c)\n",
      " Call ID: ed7edf8b-e21c-4dcf-bc53-b9b7f1ce444c\n",
      "  Args:\n",
      "    business_info: {'company_name': 'GABAS, LLC', 'description': ['tailor made AI solutions', 'consulting services'], 'industry': 'tech', 'business_type': 'software_provider', 'revenue': 4000000}\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: info_tool\n",
      "\n",
      "Error: 1 validation error for info_tool\n",
      "business_info\n",
      "  Input should be a valid dictionary or instance of BusinessInformation [type=model_type, input_value=\"{'company_name': 'GABAS,...r', 'revenue': 4000000}\", input_type=str]\n",
      "    For further information visit https://errors.pydantic.dev/2.9/v/model_type\n",
      " Please fix your mistakes.\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  info_tool (48910bb5-f882-4cd9-ba6d-4cab7a11da7e)\n",
      " Call ID: 48910bb5-f882-4cd9-ba6d-4cab7a11da7e\n",
      "  Args:\n",
      "    business_info: {'company_name': 'GABAS, LLC', 'description': ['tailor made AI solutions', 'consulting services'], 'industry': 'tech', 'business_type': 'software_provider', 'revenue': 4000000}\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: info_tool\n",
      "\n",
      "Error: 1 validation error for info_tool\n",
      "business_info\n",
      "  Input should be a valid dictionary or instance of BusinessInformation [type=model_type, input_value=\"{'company_name': 'GABAS,...r', 'revenue': 4000000}\", input_type=str]\n",
      "    For further information visit https://errors.pydantic.dev/2.9/v/model_type\n",
      " Please fix your mistakes.\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "The business information provided is:\n",
      "\n",
      "* Company Name: GABAS, LLC\n",
      "* Description: Tailor made AI solutions and consulting services\n",
      "* Industry: Tech\n",
      "* Business Type: Software Provider\n",
      "* Revenue: $4 million\n",
      "\n",
      "Please wait while I evaluate this information with the underwriters... \n",
      "\n",
      "**ToUnderwritingAssistant**\n",
      "\n",
      "Please provide a response to the business information provided.\n",
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "Looks good. \n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "Tool Calls:\n",
      "  info_tool (8a8ff5b5-0a70-41e6-91e0-d5bce6f4cb96)\n",
      " Call ID: 8a8ff5b5-0a70-41e6-91e0-d5bce6f4cb96\n",
      "  Args:\n",
      "    business_info: {'company_name': 'GABAS, LLC', 'description': ['tailor made AI solutions', 'consulting services'], 'industry': 'tech', 'business_type': 'software_provider', 'revenue': 4000000}\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: info_tool\n",
      "\n",
      "Error: 1 validation error for info_tool\n",
      "business_info\n",
      "  Input should be a valid dictionary or instance of BusinessInformation [type=model_type, input_value=\"{'company_name': 'GABAS,...r', 'revenue': 4000000}\", input_type=str]\n",
      "    For further information visit https://errors.pydantic.dev/2.9/v/model_type\n",
      " Please fix your mistakes.\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Based on the business information provided, I have gathered the following details:\n",
      "\n",
      "* Company Name: GABAS, LLC\n",
      "* Description: Tailor made AI solutions and consulting services\n",
      "* Industry: Tech\n",
      "* Business Type: Software Provider\n",
      "* Revenue: $4 million\n",
      "\n",
      "Before proceeding to quote, would you like me to provide a summary of the information provided?\n"
     ]
    }
   ],
   "source": [
    "import uuid\n",
    "from langchain.schema import AIMessage\n",
    "from langchain_core.messages import ToolMessage \n",
    "\n",
    "# Function to print the event\n",
    "def _print_event(event: dict, _printed: set, max_length=1500):\n",
    "    current_state = event.get(\"dialog_state\")\n",
    "    if current_state:\n",
    "        print(f\"Currently in: \", current_state[-1])\n",
    "    message = event.get(\"messages\")\n",
    "    if message:\n",
    "        if isinstance(message, list):\n",
    "            message = message[-1]\n",
    "        if message.id not in _printed:\n",
    "            msg_repr = message.pretty_repr(html=True)\n",
    "            if len(msg_repr) > max_length:\n",
    "                msg_repr = msg_repr[:max_length] + \" ... (truncated)\"\n",
    "            print(msg_repr)\n",
    "            _printed.add(message.id)\n",
    "\n",
    "thread_id = str(uuid.uuid4())\n",
    "\n",
    "config = {\n",
    "    \"configurable\": {\n",
    "        \"user_id\": \"1\",\n",
    "        \"thread_id\": thread_id,\n",
    "    }\n",
    "}\n",
    "\n",
    "_printed = set()\n",
    "\n",
    "# Function to process questions\n",
    "def process_questions():\n",
    "\n",
    "    global _printed  # Ensure the use of the global _printed variable\n",
    "    _printed = set()  # Re-initialize if needed for the function's scope\n",
    "    input_count = 0\n",
    "    max_inputs = 20\n",
    "\n",
    "    while input_count < max_inputs:\n",
    "        # Get user input for the question\n",
    "        question = input(\"\\nEnter your question (or type '/exit' to quit): \")\n",
    "        if question.lower() == '/exit':\n",
    "            break\n",
    "\n",
    "        # Process the question\n",
    "        events = final_graph.stream(\n",
    "            {\"messages\": (\"user\", question)}, config=config, stream_mode=\"values\"\n",
    "        )\n",
    "        for event in events:\n",
    "            _print_event(event, _printed)\n",
    "        snapshot = final_graph.get_state(config)\n",
    "\n",
    "        while snapshot.next:\n",
    "            # We have an interrupt! The agent is trying to use a tool, and the user can approve or deny it\n",
    "            user_input = input(\n",
    "                \"Do you approve of the above actions? Type 'y' to continue; \"\n",
    "                \"otherwise, explain your requested changes.\\n\\n\"\n",
    "            )\n",
    "            if user_input.strip().lower() == \"y\":\n",
    "                # Just continue\n",
    "                result = final_graph.invoke(\n",
    "                    None,\n",
    "                    config\n",
    "                )\n",
    "            else:\n",
    "                # Satisfy the tool invocation by providing instructions on the requested changes / change of mind\n",
    "                result = final_graph.invoke(\n",
    "                    {\n",
    "                        \"messages\": [\n",
    "                            ToolMessage(  # Ensure ToolMessage is imported or defined\n",
    "                                tool_call_id=event[\"messages\"][-1].tool_calls[0][\"id\"],\n",
    "                                content=f\"API call denied by user. Reasoning: '{user_input}'. Continue assisting, accounting for the user's input.\",\n",
    "                            )\n",
    "                        ]\n",
    "                    },\n",
    "                    config,\n",
    "                )\n",
    "                \n",
    "            # Filter and print the last AIMessage\n",
    "            aimessages = [message for message in result[\"messages\"] if isinstance(message, AIMessage)]\n",
    "            last_aimessage = aimessages[-1] if aimessages else None\n",
    "            if last_aimessage is not None:\n",
    "                # print(last_aimessage.content)\n",
    "                _print_event(result, _printed)\n",
    "                \n",
    "            snapshot = final_graph.get_state(config)\n",
    "\n",
    "        input_count += 1\n",
    "        if input_count >= max_inputs:\n",
    "            print(\"Maximum number of inputs reached. Terminating the loop.\")\n",
    "            break\n",
    "\n",
    "process_questions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}