Files
knightutils/test/langgraph/routing.ipynb

234 lines
26 KiB
Plaintext
Raw Normal View History

2025-10-25 21:14:41 +08:00
{
"cells": [
{
"cell_type": "code",
"id": "initial_id",
"metadata": {
"collapsed": true,
"ExecuteTime": {
"end_time": "2025-09-14T01:29:58.461138Z",
"start_time": "2025-09-14T01:29:57.277998Z"
}
},
"source": [
"import os\n",
"\n",
"os.environ['DASHSCOPE_API_KEY'] = 'sk-e2a05bbcfac84e53b73f98acef15a009'\n",
"\n",
"# Step 0: Define tools and model\n",
"\n",
"from langchain_core.tools import tool\n",
"from langchain_community.chat_models.tongyi import ChatTongyi\n",
"\n",
"llm = ChatTongyi(\n",
" model=\"qwen-max\", # 此处以qwen-max为例您可按需更换模型名称。模型列表https://help.aliyun.com/zh/model-studio/getting-started/models\n",
" streaming=True,\n",
" # other params...\n",
")"
],
"outputs": [],
"execution_count": 1
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-09-14T01:31:26.890560Z",
"start_time": "2025-09-14T01:31:26.576161Z"
}
},
"cell_type": "code",
"source": [
"from pydantic import BaseModel, Field\n",
"from typing_extensions import TypedDict\n",
"from langgraph.graph import StateGraph, START, END\n",
"from IPython.display import Image, display\n",
"from typing_extensions import Literal\n",
"from langchain_core.messages import HumanMessage, SystemMessage\n",
"\n",
"# Schema for structured output to use as routing logic\n",
"class Route(BaseModel):\n",
" step: Literal[\"poem\", \"story\", \"joke\"] = Field(\n",
" None, description=\"The next step in the routing process\"\n",
" )\n",
"\n",
"\n",
"# Augment the LLM with schema for structured output\n",
"router = llm.with_structured_output(Route)\n",
"\n",
"\n",
"# State\n",
"class State(TypedDict):\n",
" input: str\n",
" decision: str\n",
" output: str\n",
"\n",
"\n",
"# Nodes\n",
"def llm_call_1(state: State):\n",
" \"\"\"Write a story\"\"\"\n",
"\n",
" result = llm.invoke(state[\"input\"])\n",
" return {\"output\": result.content}\n",
"\n",
"\n",
"def llm_call_2(state: State):\n",
" \"\"\"Write a joke\"\"\"\n",
"\n",
" result = llm.invoke(state[\"input\"])\n",
" return {\"output\": result.content}\n",
"\n",
"\n",
"def llm_call_3(state: State):\n",
" \"\"\"Write a poem\"\"\"\n",
"\n",
" result = llm.invoke(state[\"input\"])\n",
" return {\"output\": result.content}\n",
"\n",
"\n",
"def llm_call_router(state: State):\n",
" \"\"\"Route the input to the appropriate node\"\"\"\n",
"\n",
" # Run the augmented LLM with structured output to serve as routing logic\n",
" decision = router.invoke(\n",
" [\n",
" SystemMessage(\n",
" content=\"Route the input to story, joke, or poem based on the user's request.\"\n",
" ),\n",
" HumanMessage(content=state[\"input\"]),\n",
" ]\n",
" )\n",
"\n",
" return {\"decision\": decision.step}\n",
"\n",
"\n",
"# Conditional edge function to route to the appropriate node\n",
"def route_decision(state: State):\n",
" # Return the node name you want to visit next\n",
" if state[\"decision\"] == \"story\":\n",
" return \"llm_call_1\"\n",
" elif state[\"decision\"] == \"joke\":\n",
" return \"llm_call_2\"\n",
" elif state[\"decision\"] == \"poem\":\n",
" return \"llm_call_3\"\n",
"\n",
"\n",
"# Build workflow\n",
"router_builder = StateGraph(State)\n",
"\n",
"# Add nodes\n",
"router_builder.add_node(\"llm_call_1\", llm_call_1)\n",
"router_builder.add_node(\"llm_call_2\", llm_call_2)\n",
"router_builder.add_node(\"llm_call_3\", llm_call_3)\n",
"router_builder.add_node(\"llm_call_router\", llm_call_router)\n",
"\n",
"# Add edges to connect nodes\n",
"router_builder.add_edge(START, \"llm_call_router\")\n",
"router_builder.add_conditional_edges(\n",
" \"llm_call_router\",\n",
" route_decision,\n",
" { # Name returned by route_decision : Name of next node to visit\n",
" \"llm_call_1\": \"llm_call_1\",\n",
" \"llm_call_2\": \"llm_call_2\",\n",
" \"llm_call_3\": \"llm_call_3\",\n",
" },\n",
")\n",
"router_builder.add_edge(\"llm_call_1\", END)\n",
"router_builder.add_edge(\"llm_call_2\", END)\n",
"router_builder.add_edge(\"llm_call_3\", END)\n",
"\n",
"# Compile workflow\n",
"router_workflow = router_builder.compile()"
],
"id": "3cde4ac7b799ee9c",
"outputs": [],
"execution_count": 2
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-09-14T01:31:29.539579Z",
"start_time": "2025-09-14T01:31:28.619753Z"
}
},
"cell_type": "code",
"source": [
"\n",
"# Show the workflow\n",
"display(Image(router_workflow.get_graph().draw_mermaid_png()))"
],
"id": "695f80a7d393be4e",
"outputs": [
{
"data": {
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAaMAAAFNCAIAAAAiuZdRAAAQAElEQVR4nOydB3wUxdvHZ+8ud+mNkEASIAQIJXSpUqRKE6kiTZCidKVKl6qgVJUmwp8qvTcBQeBFQAi91xBCKqTXy+X29n3uNhxHcgl3Ibdseb7ifXZna3ZmfvvM88zOKBiGIQiCIKJGQRAEQcQOKh2CIOIHlQ5BEPGDSocgiPhBpUMQRPyg0iEIIn5Q6RAzXDoWHxOWmZlG01qSrWEoGSEM/EcoijA6YlilYBnQMfoNcoWM0TE6nX5PilD6BT0EFmQy/X7sJjgB/M8YEtkUOJtMLtPRupwUimK7PcGxsAwp7P0YjjVcGm6AIcarGG9YodSfy95B5uGjqtLQpUQZR4IgJlDYnw4xcnBNVMzTTE0mo1ASO5XMTimTyQmtofSSZCgmRqVj5UafaPgnU+iXGfq1JOmlKGdnUDeGTXnjDMZfOcPQFLtM4HCd4VZeHc7y+rSGHV6vvkJmR2idjtbo1OmMjtYf7u5l92Fnz8AqLgRBUOkQlp2/hL8I1zg4y0tXcmjxubdMJiNC5ub/Jdw6l5oUl23vKGs3oKRvoANBpA0qndS5czH57K6XDq6KdgN8vP3FpggH/4h8di/Ty0/Rc1wAQSQMKp2kObA6MvJRZuPOxao18iDiZe33j7OzyNCfyhNEqqDSSZdrp+MvH0/86kdJ1P8j/4uMfKL+6odyBJEkqHQSZe+K5y8js77+QUJmzrFN0U9vZQz9GcVOigjb8YwUjtO7Yl4810hK5oA2X5QMCHZYMz2UINIDlU6K3D6fNnhWAJEebfv7ymVk7/IIgkgMVDrJsXb6E79yDnKlnEiSAbMCIx+rM9M0BJESqHTS4n5IcmYG02WEH5Ew3qWVO5dGEkRKoNJJi3MH433Lqoi06TGmdEoCnZ2dTRDJgEonLTJTdR0GlSSSx9FZfuiPWIJIBlQ6CfH3nzF29kTlwOmwDk+ePPnkk0+I9UyaNGn//v3ENpSu7BAXga46CYFKJyGiQ9XuxZWEW+7evUsKRaEPtIS6rdyy1DqCSAZUOgmRmUF7l7KV0qWmpi5YsKBTp05NmjQZMmTIvn37IHHVqlWzZs2KiYmpU6fOn3/+CSnbt28fOXJks2bN2rRpM3ny5IiInA4f27Ztg5TTp0/Xq1dv4cKFsH9UVNScOXNgT2ID3Io7yBXkwfVkgkgDVDoJodMyJUrbE9sAinbz5k0Qr127dlWtWnXevHmwOnTo0H79+pUoUeLy5ct9+vS5fv06qGGNGjVAy2D/hISEadOmsYcrlcr09HQ4dvbs2T169Dh37hwkTp8+HbSP2Aa5nHrxFBuwUgFH4pQQDEM8Stgq8Hr16lUQtQYNGsDyqFGjWrVq5e7unmufatWq7dixo3Tp0gqFvuBB9HPMmDHJyclubm4URanV6v79+9etWxc2ZWVlERsjl8vSkmmCSANUOgnB6AfDtFWH4Zo1a27evDkpKal27doNGzasXLly3n3kcjk0VxctWnT79m2w4NhEsOxA6djl4OBgwhkUtmgkBOa1hKAYJiVeTWzDzJkze/fufeHChbFjx7Zu3XrlypVarTbXPmfOnIGtVapU+eOPP0JCQpYtW5ZrB2jDEq7I1tIqZ4og0gBtOgkhk5OYZ+qg2q7EBri6ug4cOHDAgAE3btw4derU2rVrXVxc+vbta7rP3r17wfQbMWIEuwpBDPL+0GYR71JS70QtHVDpJITSURbzzCb+L/C1HT16FAKv9vb2NQ08ePDg/v37eXcrWfJ1v+V//vmHvCey1TQ05qvUdyeINMDWq4Qo7q9KjLFJtBEiDKtXr544cSIYdPHx8YcPHwaZA72DTRB/iIuLgxDqs2fPgoKC/vvvP4jDQsOW7XQCREdH5z2hSqXy9vY27kyKmnOHXlJY9qUE5raEaNateHaWTQZedXJyWrBgwYsXLwYNGtSmTZuNGzeOHj26a9eusKlx48YgeePHjz927Njw4cM//PBDcNVByCImJmbWrFngs/vmm2/AHsx7TmgLgy9v3LhxmZmZpKh5fCPdwwcbNBICxxyWFisnPC5V0fGTwb5E2qwY97jrKN8SATgtrFRAm05aVG/q9uxuBpE2e5dHyBUUypykQANeWjTqWPzW2eS/N0e37mt+RJO5c+eeOHHC7Cbwl7E9fvMyc+ZMG322BRRw5gJuafv27T4+PmY3RT5Wd/jKmyBSAluvkiP8YdqBlTEjl5ifREKtVuc3cFsBsuLg4JDfpnengM4oBdwSuA7Nzs+96YenOob0n1aWIFIClU6KQPMt8UX2wFmSq+0hf8eHHEscvhAnfpUc6KeTIl1G+CvsqD/nhxEpkZyQeekoypxEQZtOuhxcExkXlTXg+0AiAR5dTzm+6cWIRShzEgWVTtJs+jEsK1M3eI7IxW7Xr89jw7JGLEaZky6odFLn6Iaox9czSgaquo0qRUTHlVMJF48kKOxkX/8oCdMVyQ9UOoTQNL1hdnhmGu1V0u6Dtp7lq7oQ4fPX+uiwu+k6HanawOWj7j4EkTaodEgOYffT/m93XGqClqKIylHm4qFwcJYr7RU0/UYJga1QZCjDcEc5ZYfSD31n3ApBLpMjGP1mww7sgUZkFKNjqDfOA/vITFZfndi4D3tlWJbLGFpHySiiY3L2kctJtobOStelxGsz0mmGJgolKV/DuVXvEgRBUOmQvNz6N/HJrfSUhGxtFkPTRKsxo3SmC29u1esaaFje0xr3h18odaBNDKPfmWKMOmk4nOQpkvp99HsZFVUmJzqavBJQfRmWKyhQSbkdsXeS+5ZVNeuOAoe8ASodwjWnT58+ePDgokWLCIJwBX4NhnBNAR82IIiNwAKHcA0qHcI9WOAQrkGlQ7gHCxzCNdnZ2XZ2dgRBOASVDuEatOkQ7sECh3ANKh3CPVjgEK5BpUO4BwscwjXop0O4B5UO4Rq06RDuwZE4Ea5BpUO4B5UO4RpUOoR7sMAhXINKh3APFjiEazAigXAPKh3CNWjTIdyDBQ7hGlQ6hHuwwCFcg0qHcA8WOIRr0E+HcA8qHcI1aNMh3IMFDuEaVDqEe7DAIVyDSodwDxY4hGtQ6RDuwQKHcA1GJBDuQaVDuAZtOoR7sMAhXOPp6SmXywmCcAgqHcI1ycnJGo2GIAiHoNIhXANNV2jAEgThEFQ6hGtQ6RDuQaVDuAaVDuEeVDqEa1DpEO5BpUO4BpUO4R5UOoRrUOkQ7kGlQ7gGlQ7hHlQ6hGtQ6RDuQaVDuAaVDuEeVDqEa1DpEO5BpUO4BpUO4R5UOoRrUOkQ7kGlQ7gGlQ7hHlQ6hGtQ6RDukREE4RZUOoR70KZDuAaVDuEeimEYgiC2p0OHDjExMVDeKIpiU3Q6nb+//8GDBwmC2BhsvSIc0bNnT7DmZDIZ9QpIbNmyJUEQ24NKh3BE3759wYIzTSldunSPHj0IgtgeVDqEI8CI++KLL1QqlTGlbt26vr6+BEFsDyodwh2dO3c2mnXe3t69e/cmCMIJqHQIp/Tr18/R0REWatWqFRgYSBCEEzD2Kkiu/BOXEEVna83nnVxOaNo0AXajCj6h8RCKQIF4Y2fKkJQXhYzS6sxsyG9/QEYROCIkJCQjM6NmjZpubm6m6W+cxHAbZk8lk0HQ1sxF4abz3o5MTnR03kRKRzO5Hkve077azfT+GXtn0qSzN85XKzhQ6QTG42vJJ7e9hDxT2Mk0avN5J1NQOq2VGmQ8JM9OFNRvc4qmkMu0tM6qq1AywhiOMO1ror+6jNLlljr9acyWTzM7G3QKds6bLldQtDafRMqgdfmfNu+xkAJyqs0mniXsek0oQxDhgEonJJ7eSzuyNqZ+u2IV63gQ5P2xbeFjH3/7T4f4E0QgoNIJhpcv0nb8FNPv+/IE4QG7fwl1dJH3GIOWnTDAiIRgOL72pUcJ/HqPL7T8wv9lRDZBBAIqnWBIS9b5VXAkCD9w91TKFeTGvwkEEQJoIwgGrYZR2CkJwhsYHZWegEMVCANUOsHA6C1wHUF4g45mGB12NxEGqHQIgogfVDoEQcQPKh2CIOIHlQ5BCgul75B
"text/plain": [
"<IPython.core.display.Image object>"
]
},
"metadata": {},
"output_type": "display_data",
"jetTransient": {
"display_id": null
}
}
],
"execution_count": 3
},
{
"metadata": {
"ExecuteTime": {
"end_time": "2025-09-14T01:35:03.035331Z",
"start_time": "2025-09-14T01:35:00.799287Z"
}
},
"cell_type": "code",
"source": [
"\n",
"# Invoke\n",
"state = router_workflow.invoke({\"input\": \"Write me a joke about cats\"})\n",
"print(state[\"output\"])"
],
"id": "721ff52941281949",
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Why did the cat join the gym? Because it wanted to improve its mouse-cles!\n"
]
}
],
"execution_count": 4
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": "",
"id": "f1709e07d4bb0b66"
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 2
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython2",
"version": "2.7.6"
}
},
"nbformat": 4,
"nbformat_minor": 5
}