Skip to content

Commit 1fd01df

Browse files
rlancemartinrlancemartin
committed
lint
1 parent 33a55ca commit 1fd01df

File tree

7 files changed

+189
-125
lines changed

7 files changed

+189
-125
lines changed

‎examples/customer_support/src/agent/customer_support.ipynb‎

Lines changed: 89 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,12 @@
2020
"import os\n",
2121
"import getpass\n",
2222
"\n",
23-
"os.environ['ANTHROPIC_API_KEY'] = \"\"\n",
24-
"os.environ['OPENAI_API_KEY'] = \"\"\n",
23+
"os.environ[\"ANTHROPIC_API_KEY\"] = \"\"\n",
24+
"os.environ[\"OPENAI_API_KEY\"] = \"\"\n",
2525
"\n",
2626
"# turn on langsmith tracing\n",
27-
"os.environ['LANGSMITH_API_KEY'] = \"\"\n",
28-
"os.environ['LANGSMITH_TRACING_V2'] = \"true\""
27+
"os.environ[\"LANGSMITH_API_KEY\"] = \"\"\n",
28+
"os.environ[\"LANGSMITH_TRACING_V2\"] = \"true\""
2929
]
3030
},
3131
{
@@ -36,6 +36,7 @@
3636
"outputs": [],
3737
"source": [
3838
"from langchain_openai import ChatOpenAI\n",
39+
"\n",
3940
"model = ChatOpenAI(model=\"gpt-4o\")"
4041
]
4142
},
@@ -67,8 +68,24 @@
6768
"# Mock data for tools\n",
6869
"RESERVATIONS = defaultdict(lambda: {\"flight_info\": {}, \"hotel_info\": {}})\n",
6970
"TOMORROW = (datetime.date.today() + datetime.timedelta(days=1)).isoformat()\n",
70-
"FLIGHTS = [{\"departure_airport\": \"BOS\", \"arrival_airport\": \"JFK\", \"airline\": \"Jet Blue\", \"date\": TOMORROW, \"id\": \"1\"}]\n",
71-
"HOTELS = [{\"location\": \"New York\", \"name\": \"McKittrick Hotel\", \"neighborhood\": \"Chelsea\", \"id\": \"1\"}]\n",
71+
"FLIGHTS = [\n",
72+
" {\n",
73+
" \"departure_airport\": \"BOS\",\n",
74+
" \"arrival_airport\": \"JFK\",\n",
75+
" \"airline\": \"Jet Blue\",\n",
76+
" \"date\": TOMORROW,\n",
77+
" \"id\": \"1\",\n",
78+
" }\n",
79+
"]\n",
80+
"HOTELS = [\n",
81+
" {\n",
82+
" \"location\": \"New York\",\n",
83+
" \"name\": \"McKittrick Hotel\",\n",
84+
" \"neighborhood\": \"Chelsea\",\n",
85+
" \"id\": \"1\",\n",
86+
" }\n",
87+
"]\n",
88+
"\n",
7289
"\n",
7390
"# Flight tools\n",
7491
"def search_flights(\n",
@@ -86,6 +103,7 @@
86103
" # return all flights for simplicity\n",
87104
" return FLIGHTS\n",
88105
"\n",
106+
"\n",
89107
"def book_flight(\n",
90108
" flight_id: str,\n",
91109
" config: RunnableConfig,\n",
@@ -96,10 +114,9 @@
96114
" RESERVATIONS[user_id][\"flight_info\"] = flight\n",
97115
" return \"Successfully booked flight\"\n",
98116
"\n",
117+
"\n",
99118
"# Hotel tools\n",
100-
"def search_hotels(\n",
101-
" location: str\n",
102-
") -> list[dict]:\n",
119+
"def search_hotels(location: str) -> list[dict]:\n",
103120
" \"\"\"Search hotels.\n",
104121
"\n",
105122
" Args:\n",
@@ -108,6 +125,7 @@
108125
" # return all hotels for simplicity\n",
109126
" return HOTELS\n",
110127
"\n",
128+
"\n",
111129
"def book_hotel(\n",
112130
" hotel_id: str,\n",
113131
" config: RunnableConfig,\n",
@@ -118,16 +136,18 @@
118136
" RESERVATIONS[user_id][\"hotel_info\"] = hotel\n",
119137
" return \"Successfully booked hotel\"\n",
120138
"\n",
139+
"\n",
121140
"# Handoff tools\n",
122141
"transfer_to_hotel_assistant = create_handoff_tool(\n",
123142
" agent_name=\"hotel_assistant\",\n",
124-
" description=\"Transfer user to the hotel-booking assistant that can search for and book hotels.\"\n",
143+
" description=\"Transfer user to the hotel-booking assistant that can search for and book hotels.\",\n",
125144
")\n",
126145
"transfer_to_flight_assistant = create_handoff_tool(\n",
127146
" agent_name=\"flight_assistant\",\n",
128-
" description=\"Transfer user to the flight-booking assistant that can search for and book flights.\"\n",
147+
" description=\"Transfer user to the flight-booking assistant that can search for and book flights.\",\n",
129148
")\n",
130149
"\n",
150+
"\n",
131151
"# Define agent prompt\n",
132152
"def make_prompt(base_system_prompt: str) -> Callable[[dict, RunnableConfig], list]:\n",
133153
" def prompt(state: dict, config: RunnableConfig) -> list:\n",
@@ -142,26 +162,26 @@
142162
"\n",
143163
" return prompt\n",
144164
"\n",
165+
"\n",
145166
"# Define agents\n",
146167
"flight_assistant = create_react_agent(\n",
147168
" model,\n",
148169
" [search_flights, book_flight, transfer_to_hotel_assistant],\n",
149170
" prompt=make_prompt(\"You are a flight booking assistant\"),\n",
150-
" name=\"flight_assistant\"\n",
171+
" name=\"flight_assistant\",\n",
151172
")\n",
152173
"\n",
153174
"hotel_assistant = create_react_agent(\n",
154175
" model,\n",
155176
" [search_hotels, book_hotel, transfer_to_flight_assistant],\n",
156177
" prompt=make_prompt(\"You are a hotel booking assistant\"),\n",
157-
" name=\"hotel_assistant\"\n",
178+
" name=\"hotel_assistant\",\n",
158179
")\n",
159180
"\n",
160181
"# Compile and run!\n",
161182
"checkpointer = MemorySaver()\n",
162183
"builder = create_swarm(\n",
163-
" [flight_assistant, hotel_assistant],\n",
164-
" default_active_agent=\"flight_assistant\"\n",
184+
" [flight_assistant, hotel_assistant], default_active_agent=\"flight_assistant\"\n",
165185
")\n",
166186
"\n",
167187
"# Important: compile the swarm with a checkpointer to remember\n",
@@ -208,6 +228,7 @@
208228
"outputs": [],
209229
"source": [
210230
"import uuid\n",
231+
"\n",
211232
"config = {\"configurable\": {\"thread_id\": str(uuid.uuid4()), \"user_id\": \"1\"}}"
212233
]
213234
},
@@ -237,7 +258,9 @@
237258
" if isinstance(node_updates, tuple):\n",
238259
" print(node_updates)\n",
239260
" continue\n",
240-
" messages_key = next((k for k in node_updates.keys() if \"messages\" in k), None)\n",
261+
" messages_key = next(\n",
262+
" (k for k in node_updates.keys() if \"messages\" in k), None\n",
263+
" )\n",
241264
" if messages_key is not None:\n",
242265
" node_updates[messages_key][-1].pretty_print()\n",
243266
" else:\n",
@@ -314,11 +337,20 @@
314337
}
315338
],
316339
"source": [
317-
"print_stream(app.stream(\n",
318-
" {\"messages\": [{\"role\": \"user\", \"content\": \"i am looking for a flight from boston to ny tomorrow\"}]},\n",
319-
" config,\n",
320-
" subgraphs=True\n",
321-
"))"
340+
"print_stream(\n",
341+
" app.stream(\n",
342+
" {\n",
343+
" \"messages\": [\n",
344+
" {\n",
345+
" \"role\": \"user\",\n",
346+
" \"content\": \"i am looking for a flight from boston to ny tomorrow\",\n",
347+
" }\n",
348+
" ]\n",
349+
" },\n",
350+
" config,\n",
351+
" subgraphs=True,\n",
352+
" )\n",
353+
")"
322354
]
323355
},
324356
{
@@ -377,11 +409,13 @@
377409
}
378410
],
379411
"source": [
380-
"print_stream(app.stream(\n",
381-
" {\"messages\": [{\"role\": \"user\", \"content\": \"yes please\"}]},\n",
382-
" config,\n",
383-
" subgraphs=True\n",
384-
"))"
412+
"print_stream(\n",
413+
" app.stream(\n",
414+
" {\"messages\": [{\"role\": \"user\", \"content\": \"yes please\"}]},\n",
415+
" config,\n",
416+
" subgraphs=True,\n",
417+
" )\n",
418+
")"
385419
]
386420
},
387421
{
@@ -479,11 +513,17 @@
479513
}
480514
],
481515
"source": [
482-
"print_stream(app.stream(\n",
483-
" {\"messages\": [{\"role\": \"user\", \"content\": \"now i'd like to book a hotel as well\"}]},\n",
484-
" config,\n",
485-
" subgraphs=True\n",
486-
"))"
516+
"print_stream(\n",
517+
" app.stream(\n",
518+
" {\n",
519+
" \"messages\": [\n",
520+
" {\"role\": \"user\", \"content\": \"now i'd like to book a hotel as well\"}\n",
521+
" ]\n",
522+
" },\n",
523+
" config,\n",
524+
" subgraphs=True,\n",
525+
" )\n",
526+
")"
487527
]
488528
},
489529
{
@@ -542,11 +582,13 @@
542582
}
543583
],
544584
"source": [
545-
"print_stream(app.stream(\n",
546-
" {\"messages\": [{\"role\": \"user\", \"content\": \"yes please\"}]},\n",
547-
" config,\n",
548-
" subgraphs=True\n",
549-
"))"
585+
"print_stream(\n",
586+
" app.stream(\n",
587+
" {\"messages\": [{\"role\": \"user\", \"content\": \"yes please\"}]},\n",
588+
" config,\n",
589+
" subgraphs=True,\n",
590+
" )\n",
591+
")"
550592
]
551593
},
552594
{
@@ -613,11 +655,17 @@
613655
}
614656
],
615657
"source": [
616-
"print_stream(app.stream(\n",
617-
" {\"messages\": [{\"role\": \"user\", \"content\": \"i wanna talk to flight assistant now\"}]},\n",
618-
" config,\n",
619-
" subgraphs=True\n",
620-
"))"
658+
"print_stream(\n",
659+
" app.stream(\n",
660+
" {\n",
661+
" \"messages\": [\n",
662+
" {\"role\": \"user\", \"content\": \"i wanna talk to flight assistant now\"}\n",
663+
" ]\n",
664+
" },\n",
665+
" config,\n",
666+
" subgraphs=True,\n",
667+
" )\n",
668+
")"
621669
]
622670
},
623671
{

‎examples/customer_support/src/agent/customer_support.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
}
3333
]
3434

35+
3536
# Flight tools
3637
def search_flights(
3738
departure_airport: str,
@@ -124,5 +125,7 @@ def prompt(state: dict, config: RunnableConfig) -> list:
124125
)
125126

126127
# Compile and run!
127-
builder = create_swarm([flight_assistant, hotel_assistant], default_active_agent="flight_assistant")
128-
app = builder.compile()
128+
builder = create_swarm(
129+
[flight_assistant, hotel_assistant], default_active_agent="flight_assistant"
130+
)
131+
app = builder.compile()

0 commit comments

Comments
 (0)