{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Opik with LangGraph\n",
    "\n",
    "This notebook showcases how to use Opik with LangGraph. [LangGraph](https://langchain-ai.github.io/langgraph/) is a library for building stateful, multi-actor applications with LLMs, used to create agent and multi-agent workflows\n",
    "\n",
    "In this notebook, we will create a simple LangGraph workflow and focus on how to track it's execution with Opik. To learn more about LangGraph, check out the [official documentation](https://langchain-ai.github.io/langgraph/).\n",
    "\n",
    "## Creating an account on Opik Cloud\n",
    "\n",
    "[Comet](https://www.comet.com/site?from=llm&utm_source=opik&utm_medium=colab&utm_content=langgraph&utm_campaign=opik) provides a hosted version of the Opik platform, [simply create an account](https://www.comet.com/signup?from=llm&=opik&utm_medium=colab&utm_content=langgraph&utm_campaign=opik) and grab your API Key.\n",
    "\n",
    "> You can also run the Opik platform locally, see the [installation guide](https://www.comet.com/docs/opik/self-host/overview/?from=llm&utm_source=opik&utm_medium=colab&utm_content=langgraph&utm_campaign=opik) for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --quiet -U langchain langgraph opik"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import opik\n",
    "\n",
    "opik.configure(use_local=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the LangGraph graph\n",
    "\n",
    "The LangGraph graph we will be created in made up of 3 nodes:\n",
    "\n",
    "1. `classify_input`: Classify the input question\n",
    "2. `handle_greeting`: Handle the greeting question\n",
    "3. `handle_search`: Handle the search question\n",
    "\n",
    "*Note*: We will not be using any LLM calls or tools in this example to keep things simple. However in most cases, you will want to use tools to interact with external systems."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will start by creating simple functions to classify the input question and handle the greeting and search questions.\n",
    "def classify(question: str) -> str:\n",
    "    return \"greeting\" if question.startswith(\"Hello\") else \"search\"\n",
    "\n",
    "\n",
    "def classify_input_node(state):\n",
    "    question = state.get(\"question\", \"\").strip()\n",
    "    classification = classify(question)  # Assume a function that classifies the input\n",
    "    return {\"classification\": classification}\n",
    "\n",
    "\n",
    "def handle_greeting_node(state):\n",
    "    return {\"response\": \"Hello! How can I help you today?\"}\n",
    "\n",
    "\n",
    "def handle_search_node(state):\n",
    "    question = state.get(\"question\", \"\").strip()\n",
    "    search_result = f\"Search result for '{question}'\"\n",
    "    return {\"response\": search_result}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph, END\n",
    "\n",
    "from typing import TypedDict, Optional\n",
    "\n",
    "\n",
    "class GraphState(TypedDict):\n",
    "    question: Optional[str] = None\n",
    "    classification: Optional[str] = None\n",
    "    response: Optional[str] = None\n",
    "\n",
    "\n",
    "workflow = StateGraph(GraphState)\n",
    "workflow.add_node(\"classify_input\", classify_input_node)\n",
    "workflow.add_node(\"handle_greeting\", handle_greeting_node)\n",
    "workflow.add_node(\"handle_search\", handle_search_node)\n",
    "\n",
    "\n",
    "def decide_next_node(state):\n",
    "    return (\n",
    "        \"handle_greeting\"\n",
    "        if state.get(\"classification\") == \"greeting\"\n",
    "        else \"handle_search\"\n",
    "    )\n",
    "\n",
    "\n",
    "workflow.add_conditional_edges(\n",
    "    \"classify_input\",\n",
    "    decide_next_node,\n",
    "    {\"handle_greeting\": \"handle_greeting\", \"handle_search\": \"handle_search\"},\n",
    ")\n",
    "\n",
    "workflow.set_entry_point(\"classify_input\")\n",
    "workflow.add_edge(\"handle_greeting\", END)\n",
    "workflow.add_edge(\"handle_search\", END)\n",
    "\n",
    "app = workflow.compile()\n",
    "\n",
    "# Display the graph\n",
    "try:\n",
    "    from IPython.display import Image, display\n",
    "\n",
    "    display(Image(app.get_graph().draw_mermaid_png()))\n",
    "except Exception:\n",
    "    # This requires some extra dependencies and is optional\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calling the graph with Opik tracing enabled\n",
    "\n",
    "In order to log the execution of the graph, we need to define the OpikTracer callback:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from opik.integrations.langchain import OpikTracer\n",
    "\n",
    "tracer = OpikTracer(graph=app.get_graph(xray=True))\n",
    "inputs = {\"question\": \"Hello, how are you?\"}\n",
    "result = app.invoke(inputs, config={\"callbacks\": [tracer]})\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The graph execution is now logged on the Opik platform and can be viewed in the UI:\n",
    "\n",
    "![LangGraph screenshot](https://raw.githubusercontent.com/comet-ml/opik/main/apps/opik-documentation/documentation/fern/img/cookbook/langgraph_cookbook.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}