{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangGraph Integration with AgentOps\n",
    "\n",
    "This example demonstrates how to use LangGraph with AgentOps for comprehensive observability of your graph-based agent workflows.\n",
    "\n",
    "LangGraph is a framework for building stateful, multi-step applications with LLMs. AgentOps automatically instruments LangGraph to track:\n",
    "- Graph compilation and structure\n",
    "- Node executions and transitions\n",
    "- Tool usage within the graph\n",
    "- LLM calls made by agents\n",
    "- Complete execution flow with timing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install agentops langgraph langchain-openai python-dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, let's import the necessary libraries and initialize AgentOps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import Annotated, Literal, TypedDict\n",
    "from langgraph.graph import StateGraph, END\n",
    "from langgraph.graph.message import add_messages\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.messages import HumanMessage, ToolMessage\n",
    "from langchain_core.tools import tool\n",
    "import agentops\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# Load environment variables\n",
    "load_dotenv()\n",
    "\n",
    "# Initialize AgentOps - this enables automatic instrumentation\n",
    "agentops.init(os.getenv(\"AGENTOPS_API_KEY\"), auto_start_session=False)\n",
    "trace = agentops.start_trace(\"langgraph_example\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Tools\n",
    "\n",
    "Let's create some simple tools that our agent can use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "def get_weather(location: str) -> str:\n",
    "    \"\"\"Get the weather for a given location.\"\"\"\n",
    "    # Simulated weather data\n",
    "    weather_data = {\n",
    "        \"New York\": \"Sunny, 72°F\",\n",
    "        \"London\": \"Cloudy, 60°F\",\n",
    "        \"Tokyo\": \"Rainy, 65°F\",\n",
    "        \"Paris\": \"Partly cloudy, 68°F\",\n",
    "        \"Sydney\": \"Clear, 75°F\",\n",
    "    }\n",
    "    return weather_data.get(location, f\"Weather data not available for {location}\")\n",
    "\n",
    "\n",
    "@tool\n",
    "def calculate(expression: str) -> str:\n",
    "    \"\"\"Evaluate a mathematical expression.\"\"\"\n",
    "    try:\n",
    "        result = eval(expression)\n",
    "        return f\"The result is: {result}\"\n",
    "    except Exception as e:\n",
    "        return f\"Error calculating expression: {str(e)}\"\n",
    "\n",
    "\n",
    "# Collect tools for binding to the model\n",
    "tools = [get_weather, calculate]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define Agent State\n",
    "\n",
    "In LangGraph, we need to define the state that will be passed between nodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AgentState(TypedDict):\n",
    "    messages: Annotated[list, add_messages]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the Model and Node Functions\n",
    "\n",
    "We'll create a model with tool binding and define the functions that will be our graph nodes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create model with tool binding\n",
    "model = ChatOpenAI(temperature=0, model=\"gpt-4o-mini\").bind_tools(tools)\n",
    "\n",
    "\n",
    "def should_continue(state: AgentState) -> Literal[\"tools\", \"end\"]:\n",
    "    \"\"\"Determine if we should continue to tools or end.\"\"\"\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "\n",
    "    # If the LLM wants to use tools, continue to the tools node\n",
    "    if hasattr(last_message, \"tool_calls\") and last_message.tool_calls:\n",
    "        return \"tools\"\n",
    "    # Otherwise, we're done\n",
    "    return \"end\"\n",
    "\n",
    "\n",
    "def call_model(state: AgentState):\n",
    "    \"\"\"Call the language model.\"\"\"\n",
    "    messages = state[\"messages\"]\n",
    "    response = model.invoke(messages)\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "def call_tools(state: AgentState):\n",
    "    \"\"\"Execute the tool calls requested by the model.\"\"\"\n",
    "    messages = state[\"messages\"]\n",
    "    last_message = messages[-1]\n",
    "\n",
    "    tool_messages = []\n",
    "    for tool_call in last_message.tool_calls:\n",
    "        tool_name = tool_call[\"name\"]\n",
    "        tool_args = tool_call[\"args\"]\n",
    "\n",
    "        # Find and execute the requested tool\n",
    "        for available_tool in tools:\n",
    "            if available_tool.name == tool_name:\n",
    "                result = available_tool.invoke(tool_args)\n",
    "                tool_messages.append(ToolMessage(content=str(result), tool_call_id=tool_call[\"id\"]))\n",
    "                break\n",
    "\n",
    "    return {\"messages\": tool_messages}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Graph\n",
    "\n",
    "Now let's construct the LangGraph workflow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the graph\n",
    "workflow = StateGraph(AgentState)\n",
    "\n",
    "# Add nodes\n",
    "workflow.add_node(\"agent\", call_model)\n",
    "workflow.add_node(\"tools\", call_tools)\n",
    "\n",
    "# Set the entry point\n",
    "workflow.set_entry_point(\"agent\")\n",
    "\n",
    "# Add conditional edges\n",
    "workflow.add_conditional_edges(\"agent\", should_continue, {\"tools\": \"tools\", \"end\": END})\n",
    "\n",
    "# Add edge from tools back to agent\n",
    "workflow.add_edge(\"tools\", \"agent\")\n",
    "\n",
    "# Compile the graph\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Examples\n",
    "\n",
    "Let's test our agent with different queries that require tool usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example 1: Weather query\n",
    "print(\"Example 1: Weather Query\")\n",
    "print(\"=\" * 50)\n",
    "\n",
    "messages = [HumanMessage(content=\"What's the weather in New York and Tokyo?\")]\n",
    "result = app.invoke({\"messages\": messages})\n",
    "\n",
    "final_message = result[\"messages\"][-1]\n",
    "print(f\"Response: {final_message.content}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example 2: Math calculation\n",
    "print(\"Example 2: Math Calculation\")\n",
    "print(\"=\" * 50)\n",
    "\n",
    "messages = [HumanMessage(content=\"Calculate 25 * 4 + 10\")]\n",
    "result = app.invoke({\"messages\": messages})\n",
    "\n",
    "final_message = result[\"messages\"][-1]\n",
    "print(f\"Response: {final_message.content}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example 3: Combined query\n",
    "print(\"Example 3: Combined Query\")\n",
    "print(\"=\" * 50)\n",
    "\n",
    "messages = [HumanMessage(content=\"What's the weather in Paris? Also calculate 100/5\")]\n",
    "result = app.invoke({\"messages\": messages})\n",
    "\n",
    "final_message = result[\"messages\"][-1]\n",
    "print(f\"Response: {final_message.content}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## View in AgentOps Dashboard\n",
    "\n",
    "After running this notebook, you can view the traces in your AgentOps dashboard. You'll see:\n",
    "\n",
    "1. **Graph Compilation**: The structure of your LangGraph with nodes and edges\n",
    "2. **Execution Flow**: How the graph executed, including:\n",
    "   - Agent node calls\n",
    "   - Tool node executions\n",
    "   - State transitions\n",
    "3. **LLM Calls**: Each ChatGPT call with prompts and completions\n",
    "4. **Tool Usage**: Which tools were called and their results\n",
    "5. **Timing Information**: How long each step took\n",
    "\n",
    "The instrumentation captures the full context of your LangGraph application automatically!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"✅ Check your AgentOps dashboard for comprehensive traces!\")\n",
    "print(\"🔍 You'll see the graph structure, execution flow, and all LLM/tool calls.\")\n",
    "agentops.end_trace(trace)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}