From 867496c7e392abe0514c806514b4d55ed9ad0780 Mon Sep 17 00:00:00 2001
From: bracesproul <braceasproul@gmail.com>
Date: Mon, 19 Aug 2024 16:37:52 -0700
Subject: [PATCH] updated streaming tokens without langchain doc

---
 .../streaming-tokens-without-langchain.ipynb  | 50 +++++++++----------
 1 file changed, 24 insertions(+), 26 deletions(-)

diff --git a/examples/how-tos/streaming-tokens-without-langchain.ipynb b/examples/how-tos/streaming-tokens-without-langchain.ipynb
index 309fed9b8..93955eb54 100644
--- a/examples/how-tos/streaming-tokens-without-langchain.ipynb
+++ b/examples/how-tos/streaming-tokens-without-langchain.ipynb
@@ -79,17 +79,21 @@
    "source": [
     "import { dispatchCustomEvent } from \"@langchain/core/callbacks/dispatch\";\n",
     "import { wrapOpenAI } from \"langsmith/wrappers/openai\";\n",
+    "import { Annotation } from \"@langchain/langgraph\";\n",
     "\n",
-    "type GraphState = {\n",
-    "  messages: OpenAI.ChatCompletionMessageParam[];\n",
-    "};\n",
+    "const GraphState = Annotation.Root({\n",
+    "  messages: Annotation<OpenAI.ChatCompletionMessageParam[]>({\n",
+    "    reducer: (x, y) => x.concat(y),\n",
+    "  }),\n",
+    "});\n",
     "\n",
     "// If using LangSmith, use \"wrapOpenAI\" on the whole client or\n",
     "// \"traceable\" to wrap a single method for nicer tracing:\n",
     "// https://docs.smith.langchain.com/how_to_guides/tracing/annotate_code\n",
     "const wrappedClient = wrapOpenAI(openaiClient);\n",
     "\n",
-    "const callModel = async ({ messages }: GraphState) => {\n",
+    "const callModel = async (state: typeof GraphState.State): Promise<Partial<typeof GraphState.State>> => {\n",
+    "  const { messages } = state;\n",
     "  const stream = await wrappedClient.chat.completions.create({\n",
     "    messages,\n",
     "    model: \"gpt-4o-mini\",\n",
@@ -97,7 +101,7 @@
     "    stream: true,\n",
     "  });\n",
     "  let responseContent = \"\";\n",
-    "  let role = \"assistant\";\n",
+    "  let role: string = \"assistant\";\n",
     "  let toolCallId: string | undefined;\n",
     "  let toolCallName: string | undefined;\n",
     "  let toolCallArgs = \"\";\n",
@@ -133,11 +137,12 @@
     "        name: toolCallName,\n",
     "        arguments: toolCallArgs\n",
     "      },\n",
-    "      type: \"function\",\n",
+    "      type: \"function\" as const,\n",
     "    }];\n",
     "  }\n",
+    "\n",
     "  const responseMessage = {\n",
-    "    role,\n",
+    "    role: role as any,\n",
     "    content: responseContent,\n",
     "    tool_calls: finalToolCalls,\n",
     "  };\n",
@@ -172,7 +177,8 @@
     "  }\n",
     "};\n",
     "\n",
-    "const callTools = async ({ messages }: GraphState) => {\n",
+    "const callTools = async (state: typeof GraphState.State): Promise<Partial<typeof GraphState.State>> => {\n",
+    "  const { messages } = state;\n",
     "  const mostRecentMessage = messages[messages.length - 1];\n",
     "  const toolCalls = (mostRecentMessage as OpenAI.ChatCompletionAssistantMessageParam).tool_calls;\n",
     "  if (toolCalls === undefined || toolCalls.length === 0) {\n",
@@ -186,7 +192,7 @@
     "  const response = await toolNameMap[functionName](functionArguments);\n",
     "  const toolMessage = {\n",
     "    tool_call_id: toolCalls[0].id,\n",
-    "    role: \"tool\",\n",
+    "    role: \"tool\" as const,\n",
     "    name: functionName,\n",
     "    content: response,\n",
     "  }\n",
@@ -210,13 +216,11 @@
    "outputs": [],
    "source": [
     "import { StateGraph } from \"@langchain/langgraph\";\n",
-    "\n",
     "import OpenAI from \"openai\";\n",
-    "type GraphState = {\n",
-    "  messages: OpenAI.ChatCompletionMessageParam[];\n",
-    "};\n",
     "\n",
-    "const shouldContinue = ({ messages }: GraphState) => {\n",
+    "// We can reuse the same `GraphState` from above as it has not changed.\n",
+    "const shouldContinue = (state: typeof GraphState.State) => {\n",
+    "  const { messages } = state;\n",
     "  const lastMessage =\n",
     "    messages[messages.length - 1] as OpenAI.ChatCompletionAssistantMessageParam;\n",
     "  if (lastMessage?.tool_calls !== undefined && lastMessage?.tool_calls.length > 0) {\n",
@@ -225,15 +229,7 @@
     "  return \"__end__\";\n",
     "}\n",
     "\n",
-    "const workflow = new StateGraph<GraphState>({\n",
-    "  channels: {\n",
-    "    messages: {\n",
-    "      reducer: (a: any, b: any) => a.concat(b),\n",
-    "    }\n",
-    "  }\n",
-    "});\n",
-    "\n",
-    "const graph = workflow\n",
+    "const graph = new StateGraph(GraphState)\n",
     "  .addNode(\"model\", callModel)\n",
     "  .addNode(\"tools\", callTools)\n",
     "  .addEdge(\"__start__\", \"model\")\n",
@@ -242,7 +238,8 @@
     "    __end__: \"__end__\",\n",
     "  })\n",
     "  .addEdge(\"tools\", \"model\")\n",
-    "  .compile();"
+    "  .compile();\n",
+    "  "
    ]
   },
   {
@@ -288,7 +285,7 @@
      "text": [
       "streamed_tool_call_chunk {\n",
       "  index: 0,\n",
-      "  id: 'call_1Lt33wOgPfjXwvM9bcWoe1yZ',\n",
+      "  id: 'call_v99reml4gZvvUypPgOpLgxM2',\n",
       "  type: 'function',\n",
       "  function: { name: 'get_items', arguments: '' }\n",
       "}\n",
@@ -302,7 +299,8 @@
       "streamed_token { content: ' the' }\n",
       "streamed_token { content: ' bedroom' }\n",
       "streamed_token { content: ',' }\n",
-      "streamed_token { content: \" you'll\" }\n",
+      "streamed_token { content: ' you' }\n",
+      "streamed_token { content: ' can' }\n",
       "streamed_token { content: ' find' }\n",
       "streamed_token { content: ' socks' }\n",
       "streamed_token { content: ',' }\n",