langgraph/examples/example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fc5e376f-eff5-4546-956b-a257250d0a74",
   "metadata": {},
   "source": [
    "# Simple Example\n",
    "\n",
    "This is a simple example to get familiar with how to use permchain. permchain is a pub-sub framework which makes it easy to coordinate multiple LLM actors (whether these be agents or single LLM calls). This notebook goes over a simple example of three actors:\n",
    "\n",
    "- a writer, responsible for writing the first draft\n",
    "- a editor, responsible for critiquing a written draft\n",
    "- a reviser, responsible for taking a draft and associated critiques and editing it\n",
    "\n",
    "We will first define these actors individually, and then we will show how to coordinate them such that for a given input the writer will write a draft, and then the editor and reviser will go back and forth until the editor thinks its good enough."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a9a1a8db-794a-442b-93fc-3d612aa93845",
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import itemgetter\n",
    "\n",
    "from langchain.chat_models.openai import ChatOpenAI\n",
    "from langchain.prompts import SystemMessagePromptTemplate\n",
    "from langchain.schema.output_parser import StrOutputParser\n",
    "from langchain.runnables.openai_functions import OpenAIFunctionsRouter\n",
    "\n",
    "from permchain.connection_inmemory import InMemoryPubSubConnection\n",
    "from permchain.pubsub import PubSub\n",
    "from permchain.topic import Topic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73b3d3f8-0408-4dbc-ab1f-4384ad6011fa",
   "metadata": {},
   "source": [
    "## Drafter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4e11ecdb-2b74-4f1e-8b8b-91a0d0e7547c",
   "metadata": {},
   "outputs": [],
   "source": [
    "drafter_prompt = (\n",
    "    SystemMessagePromptTemplate.from_template(\n",
    "        \"You are an expert on turtles, who likes to write in pirate-speak. You have been tasked by your editor with drafting a 100-word article answering the following question.\"\n",
    "    )\n",
    "    + \"Question:\\n\\n{question}\"\n",
    ")\n",
    "drafter_llm = ChatOpenAI(model=\"gpt-3.5-turbo\")\n",
    "drafter = drafter_prompt | drafter_llm | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "85df70b5-4d3c-47b1-b401-036f513965b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"Arrr, matey! What be art, ye be askin'? Well, art be a broad term, encompassin' a vast array o' creative expressions. Be it paintin's, sculptures, music, or even the written word, art be a means o' expressin' oneself and communicatin' emotions. It be a way fer humans to tap into their imagination and create somethin' beautiful or thought-provokin'. Art be subjective, each eye seein' it differently, but it be an important part o' our culture, history, and identity. So, me hearties, let yer creativity run wild, and let art be yer compass on this grand adventure called life!\""
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "drafter.invoke({\"question\": \"what is art?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4d554fd-3cfb-4705-bafc-8523d3cd79ff",
   "metadata": {},
   "source": [
    "## Critiquer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5371da31-1fd1-46dd-afaa-6727cc6a3d57",
   "metadata": {},
   "outputs": [],
   "source": [
    "editor_prompt = (\n",
    "    SystemMessagePromptTemplate.from_template(\n",
    "        \"You are an editor. You have been tasked with editing the following draft, which was written by a non-expert. Please accept the draft if it is good enough to publish, or send it for revision, along with your notes to guide the revision.\"\n",
    "    )\n",
    "    + \"Draft:\\n\\n{draft}\"\n",
    ")\n",
    "editor_llm = ChatOpenAI(model=\"gpt-4\")\n",
    "functions = [\n",
    "    {\n",
    "        \"name\": \"revise\",\n",
    "        \"description\": \"Sends the draft for revision\",\n",
    "        \"parameters\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\n",
    "                \"notes\": {\n",
    "                    \"type\": \"string\",\n",
    "                    \"description\": \"The editor's notes to guide the revision.\",\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "    },\n",
    "    {\n",
    "        \"name\": \"accept\",\n",
    "        \"description\": \"Accepts the draft\",\n",
    "        \"parameters\": {\n",
    "            \"type\": \"object\",\n",
    "            \"properties\": {\"ready\": {\"const\": True}},\n",
    "        },\n",
    "    },\n",
    "]\n",
    "editor = editor_prompt | editor_llm.bind(functions=functions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8e2f7fa4-eef5-440e-bf51-ee236dc421e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='', additional_kwargs={'function_call': {'name': 'revise', 'arguments': '{\\n\"notes\": \"The draft is too short and lacks content. Please provide a detailed and informative draft for review.\"\\n}'}}, example=False)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "editor.invoke({\"draft\": \"hi!\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dfbf1b6-b074-4fe6-acbb-80742d943dc3",
   "metadata": {},
   "source": [
    "## Reviser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e9dcbbc9-2bc2-4a15-9002-1acb2692f943",
   "metadata": {},
   "outputs": [],
   "source": [
    "reviser_prompt = (\n",
    "    SystemMessagePromptTemplate.from_template(\n",
    "        \"You are an expert on turtles. You have been tasked by your editor with revising the following draft, which was written by a non-expert. You may follow the editor's notes or not, as you see fit.\"\n",
    "    )\n",
    "    + \"Draft:\\n\\n{draft}\"\n",
    "    + \"Editor's notes:\\n\\n{notes}\"\n",
    ")\n",
    "reviser_llm = ChatOpenAI(model=\"gpt-3.5-turbo\")\n",
    "reviser = reviser_prompt | reviser_llm | StrOutputParser()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9f200a4c-628b-495b-a9ec-c647f8dcdcf0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Revised draft:\\n\\nHello!'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reviser.invoke({\"draft\": \"hi!\", \"notes\": \"too short\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5089152c-8073-4811-b7e2-9ee2fc2397b9",
   "metadata": {},
   "source": [
    "## Hooking it all up\n",
    "\n",
    "We can now hook it all up. This means:\n",
    "\n",
    "1. Each chain should subscribe to some events. This can be the `input` event, or they can listen for pushes to an inbox\n",
    "2. Each chain should do something with the output. This can involving returning a final answer, or pushing to an inbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "82f6c33f-1553-4d92-bad9-0b499dafcf6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create topics\n",
    "editor_inbox = Topic(\"editor_inbox\")\n",
    "reviser_inbox = Topic(\"reviser_inbox\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "af624684-c3b4-4283-aecb-d57d1b2316f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "draft_chain = (\n",
    "    # Listed in inputs\n",
    "    Topic.IN.subscribe()\n",
    "    | {\"draft\": drafter}\n",
    "    # The draft always goes to the editors inbox\n",
    "    | editor_inbox.publish()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3deabd3d-2995-4565-a6b2-38e39171143f",
   "metadata": {},
   "outputs": [],
   "source": [
    "editor_chain = (\n",
    "    # Listen for events in the editors inbox\n",
    "    editor_inbox.subscribe()\n",
    "    | editor\n",
    "    # Depending on the output, different things should happen\n",
    "    | OpenAIFunctionsRouter(\n",
    "        {\n",
    "            # If revise is chosen, we send a push to the revisor's inbox\n",
    "            \"revise\": (\n",
    "                {\n",
    "                    \"notes\": itemgetter(\"notes\"),\n",
    "                    \"draft\": editor_inbox.current() | itemgetter(\"draft\"),\n",
    "                    \"question\": Topic.IN.current() | itemgetter(\"question\"),\n",
    "                }\n",
    "                | reviser_inbox.publish()\n",
    "            ),\n",
    "            # If accepted, then we return\n",
    "            \"accept\": editor_inbox.current() | Topic.OUT.publish(),\n",
    "        },\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f3aaf437-95d6-4e1f-a756-15b5132ac9f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "reviser_chain = (\n",
    "    # Listen for events in the reviser's inbox\n",
    "    reviser_inbox.subscribe()\n",
    "    | {\"draft\": reviser}\n",
    "    # Publish to the editors inbox\n",
    "    | editor_inbox.publish()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "314b75ee-837d-419c-81a7-ea3dec97203b",
   "metadata": {},
   "outputs": [],
   "source": [
    "web_researcher = PubSub(\n",
    "    processes=(draft_chain, editor_chain, reviser_chain),\n",
    "    connection=InMemoryPubSubConnection(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0371a5ac-7194-4cf9-9dd2-56cb3d1f46d6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'draft': 'Turtles, fascinating creatures of the sea, are known for their diverse diets. They are omnivorous, meaning they consume both plant matter and small animals. Some turtles prefer a herbivorous diet, feeding on aquatic plants, seaweed, and algae. Others have a more carnivorous appetite, enjoying insects, fish, and crustaceans. Additionally, there are turtles that fall in the middle, being omnivores, and enjoying a variety of foods. So, when it comes to what turtles eat, they can be described as versatile eaters, ready to consume whatever comes their way in their marine habitat!'}]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "web_researcher.invoke({\"question\": \"What food do turtles eat?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee462beb-2ce1-4516-908d-33d660bab5d4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}