langraph/LanggraphAgentDemo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import functools\n",
    "from pprint import pprint\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "from langchain_openai import AzureChatOpenAI\n",
    "from langchain_community.vectorstores.azuresearch import AzureSearch\n",
    "from langchain_openai import AzureOpenAIEmbeddings\n",
    "from langgraph.prebuilt import ToolNode\n",
    "from langgraph.graph import END, StateGraph\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "from utils.prompts import (\n",
    "    default_router_prompt,\n",
    "    default_chat_prompt,\n",
    "    default_tool_using_prompt,\n",
    ")\n",
    "from utils.router import AgentState, classifier_router\n",
    "from utils.agent import create_agent, agent_node, invoke_graph\n",
    "from utils.tools import search_web_tool, create_vector_store_tool\n",
    "\n",
    "# Load getenvment variables\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set up Clients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt4o_client = AzureChatOpenAI(\n",
    "    model=os.getenv(\"AZURE_OPENAI_MODEL\"),\n",
    "    azure_deployment=os.getenv(\"AZURE_OPENAI_DEPLOYMENT_NAME\"),\n",
    "    api_key=os.getenv(\"AZURE_OPENAI_API_KEY\"),\n",
    "    api_version=os.getenv(\"AZURE_OPENAI_VERSION\"),\n",
    "    azure_endpoint=os.getenv(\"AZURE_OPENAI_GPOT4O_ENDPOINT\"),\n",
    "    temperature=os.getenv(\"AZURE_OPENAI_TEMPERATURE\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings: AzureOpenAIEmbeddings = AzureOpenAIEmbeddings(\n",
    "    azure_deployment=os.getenv(\"AZURE_OPENAI_EMBEDDER_DEPLOYMENT\"),\n",
    "    openai_api_version=os.getenv(\"AZURE_OPENAI_VERSION\"),\n",
    "    azure_endpoint=os.getenv(\"AZURE_OPENAI_EMBEDDER_ENDPOINT\"),\n",
    "    api_key=os.getenv(\"AZURE_OPENAI_API_KEY\"),\n",
    "    dimensions=os.getenv(\"AZURE_OPENAI_EMBEDDER_DIMENSIONS\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store: AzureSearch = AzureSearch(\n",
    "    azure_search_endpoint=os.getenv(\"AI_SEARCH_ENDPOINT\"),\n",
    "    azure_search_key=os.getenv(\"AI_SEARCH_API_KEY\"),\n",
    "    index_name=os.getenv(\"INDEX_NAME\"),\n",
    "    embedding_function=embeddings.embed_query,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create retrieval tool\n",
    "retriever_tool = create_vector_store_tool(vector_store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier_agent = create_agent(gpt4o_client, default_router_prompt)\n",
    "chat_agent = create_agent(gpt4o_client, default_chat_prompt)\n",
    "tool_using_agent = create_agent(\n",
    "    gpt4o_client, default_tool_using_prompt, tools=[search_web_tool, retriever_tool]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "classifier_node = functools.partial(\n",
    "    agent_node, agent=classifier_agent, name=\"classifier\"\n",
    ")\n",
    "chat_node = functools.partial(agent_node, agent=chat_agent, name=\"chat\")\n",
    "tool_using_node = functools.partial(\n",
    "    agent_node, agent=tool_using_agent, name=\"tool_using\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a tool node\n",
    "tools = [search_web_tool, retriever_tool]\n",
    "tool_node = ToolNode(tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Graph Workflow\n",
    "workflow = StateGraph(AgentState)\n",
    "\n",
    "# Add Nodes to the Workflow\n",
    "workflow.add_node(\"classifier\", classifier_node)\n",
    "workflow.add_node(\"chat\", chat_node)\n",
    "workflow.add_node(\"tool_using\", tool_using_node)\n",
    "workflow.add_node(\"tool\", tool_node)\n",
    "\n",
    "# Define Edges\n",
    "workflow.add_conditional_edges(\n",
    "    \"classifier\",\n",
    "    classifier_router,\n",
    "    {\"SPECIFIC\": \"tool_using\", \"GENERAL\": \"chat\"},\n",
    ")\n",
    "\n",
    "workflow.add_edge(\"tool_using\", \"tool\")\n",
    "workflow.add_edge(\"tool\", \"chat\")\n",
    "workflow.add_edge(\"chat\", END)\n",
    "\n",
    "# Set Entry Point\n",
    "workflow.set_entry_point(\"classifier\")\n",
    "\n",
    "# Compile the Graph\n",
    "graph = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize memory to persist memory between graph runs\n",
    "checkpointer = MemorySaver()\n",
    "graph = workflow.compile(debug=True, checkpointer=checkpointer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = invoke_graph(graph=graph, messages=\"Which country started WWII?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = invoke_graph(graph=graph, messages=\"What are the main risks of Tesla?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(result)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "agent_demo",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}