DuyguJones commited on
Commit
54741ab
Β·
1 Parent(s): a0958e5
test.ipynb β†’ agents/__init__.py RENAMED
File without changes
test_gaia.ipynb ADDED
@@ -0,0 +1,584 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "cells": [
3
+ {
4
+ "cell_type": "markdown",
5
+ "id": "9eedd8e4",
6
+ "metadata": {},
7
+ "source": [
8
+ "# Testing the Agent"
9
+ ]
10
+ },
11
+ {
12
+ "cell_type": "markdown",
13
+ "id": "f906ea45",
14
+ "metadata": {},
15
+ "source": [
16
+ "## Set the Env"
17
+ ]
18
+ },
19
+ {
20
+ "cell_type": "code",
21
+ "execution_count": 3,
22
+ "id": "1f26b276",
23
+ "metadata": {},
24
+ "outputs": [],
25
+ "source": [
26
+ "# Enable auto-reloading of external modules when they change\n",
27
+ "%load_ext autoreload\n",
28
+ "\n",
29
+ "# Set auto-reload mode to 2: reload all modules (except those excluded) before executing a line\n",
30
+ "%autoreload 2"
31
+ ]
32
+ },
33
+ {
34
+ "cell_type": "code",
35
+ "execution_count": 19,
36
+ "id": "28e55a4e",
37
+ "metadata": {},
38
+ "outputs": [],
39
+ "source": [
40
+ "import os\n",
41
+ "from dotenv import load_dotenv\n",
42
+ "load_dotenv()\n",
43
+ "\n",
44
+ "## Langsmith Tracking\n",
45
+ "os.environ[\"LANGCHAIN_API_KEY\"]=os.getenv(\"LANGCHAIN_API_KEY\")\n",
46
+ "os.environ[\"LANGCHAIN_TRACING_V2\"]=\"true\"\n",
47
+ "os.environ[\"LANGCHAIN_PROJECT\"]=os.getenv(\"LANGCHAIN_PROJECT\")"
48
+ ]
49
+ },
50
+ {
51
+ "cell_type": "markdown",
52
+ "id": "2ce694a1",
53
+ "metadata": {},
54
+ "source": [
55
+ "## Load Data"
56
+ ]
57
+ },
58
+ {
59
+ "cell_type": "code",
60
+ "execution_count": null,
61
+ "id": "1a66ce56",
62
+ "metadata": {},
63
+ "outputs": [],
64
+ "source": [
65
+ "import json\n",
66
+ "import os\n",
67
+ "\n",
68
+ "# Load the metadata.jsonl file\n",
69
+ "with open('data/metadata.jsonl', 'r') as f:\n",
70
+ " json_list = list(f)\n",
71
+ "\n",
72
+ "json_QA = []\n",
73
+ "for json_str in json_list:\n",
74
+ " json_data = json.loads(json_str)\n",
75
+ " json_QA.append(json_data) "
76
+ ]
77
+ },
78
+ {
79
+ "cell_type": "markdown",
80
+ "id": "7b11f6ca",
81
+ "metadata": {},
82
+ "source": [
83
+ "### Data Analyses\n",
84
+ "\n",
85
+ "Here we analyze the data in order to find what tools we need to build a robust agent"
86
+ ]
87
+ },
88
+ {
89
+ "cell_type": "code",
90
+ "execution_count": null,
91
+ "id": "ef2da36e",
92
+ "metadata": {},
93
+ "outputs": [
94
+ {
95
+ "data": {
96
+ "text/plain": [
97
+ "{'task_id': 'c61d22de-5f6c-4958-a7f6-5e9707bd3466',\n",
98
+ " 'Question': 'A paper about AI regulation that was originally submitted to arXiv.org in June 2022 shows a figure with three axes, where each axis has a label word at both ends. Which of these words is used to describe a type of society in a Physics and Society article submitted to arXiv.org on August 11, 2016?',\n",
99
+ " 'Level': 2,\n",
100
+ " 'Final answer': 'egalitarian',\n",
101
+ " 'file_name': '',\n",
102
+ " 'Annotator Metadata': {'Steps': '1. Go to arxiv.org and navigate to the Advanced Search page.\\n2. Enter \"AI regulation\" in the search box and select \"All fields\" from the dropdown.\\n3. Enter 2022-06-01 and 2022-07-01 into the date inputs, select \"Submission date (original)\", and submit the search.\\n4. Go through the search results to find the article that has a figure with three axes and labels on each end of the axes, titled \"Fairness in Agreement With European Values: An Interdisciplinary Perspective on AI Regulation\".\\n5. Note the six words used as labels: deontological, egalitarian, localized, standardized, utilitarian, and consequential.\\n6. Go back to arxiv.org\\n7. Find \"Physics and Society\" and go to the page for the \"Physics and Society\" category.\\n8. Note that the tag for this category is \"physics.soc-ph\".\\n9. Go to the Advanced Search page.\\n10. Enter \"physics.soc-ph\" in the search box and select \"All fields\" from the dropdown.\\n11. Enter 2016-08-11 and 2016-08-12 into the date inputs, select \"Submission date (original)\", and submit the search.\\n12. Search for instances of the six words in the results to find the paper titled \"Phase transition from egalitarian to hierarchical societies driven by competition between cognitive and social constraints\", indicating that \"egalitarian\" is the correct answer.',\n",
103
+ " 'Number of steps': '12',\n",
104
+ " 'How long did this take?': '8 minutes',\n",
105
+ " 'Tools': '1. Web browser\\n2. Image recognition tools (to identify and parse a figure with three axes)',\n",
106
+ " 'Number of tools': '2'}}"
107
+ ]
108
+ },
109
+ "execution_count": 7,
110
+ "metadata": {},
111
+ "output_type": "execute_result"
112
+ }
113
+ ],
114
+ "source": [
115
+ "json_QA[0] # Display the first entry to check the structure"
116
+ ]
117
+ },
118
+ {
119
+ "cell_type": "code",
120
+ "execution_count": null,
121
+ "id": "597d7911",
122
+ "metadata": {},
123
+ "outputs": [
124
+ {
125
+ "data": {
126
+ "text/plain": [
127
+ "'A paper about AI regulation that was originally submitted to arXiv.org in June 2022 shows a figure with three axes, where each axis has a label word at both ends. Which of these words is used to describe a type of society in a Physics and Society article submitted to arXiv.org on August 11, 2016?'"
128
+ ]
129
+ },
130
+ "execution_count": 15,
131
+ "metadata": {},
132
+ "output_type": "execute_result"
133
+ }
134
+ ],
135
+ "source": [
136
+ "json_QA[0]['Question'] # Access the 'Question' field"
137
+ ]
138
+ },
139
+ {
140
+ "cell_type": "code",
141
+ "execution_count": null,
142
+ "id": "53f27fdd",
143
+ "metadata": {},
144
+ "outputs": [
145
+ {
146
+ "name": "stdout",
147
+ "output_type": "stream",
148
+ "text": [
149
+ "==================================================\n",
150
+ "Task ID: e1fc63a2-da7a-432f-be78-7c4a95598703\n",
151
+ "Question: If Eliud Kipchoge could maintain his record-making marathon pace indefinitely, how many thousand hours would it take him to run the distance between the Earth and the Moon its closest approach? Please use the minimum perigee value on the Wikipedia page for the Moon when carrying out your calculation. Round your result to the nearest 1000 hours and do not use any comma separators if necessary.\n",
152
+ "Level: 1\n",
153
+ "Final Answer: 17\n",
154
+ "Annotator Metadata: \n",
155
+ " β”œβ”€β”€ Steps: \n",
156
+ " β”‚ β”œβ”€β”€ 1. Googled Eliud Kipchoge marathon pace to find 4min 37sec/mile\n",
157
+ " β”‚ β”œβ”€β”€ 2. Converted into fractions of hours.\n",
158
+ " β”‚ β”œβ”€β”€ 3. Found moon periapsis in miles (225,623 miles).\n",
159
+ " β”‚ β”œβ”€β”€ 4. Multiplied the two to find the number of hours and rounded to the nearest 100 hours.\n",
160
+ " β”œβ”€β”€ Number of steps: 4\n",
161
+ " β”œβ”€β”€ How long did this take?: 20 Minutes\n",
162
+ " β”œβ”€β”€ Tools:\n",
163
+ " β”‚ β”œβ”€β”€ 1. A web browser.\n",
164
+ " β”‚ β”œβ”€β”€ 2. A search engine.\n",
165
+ " β”‚ β”œβ”€β”€ 3. A calculator.\n",
166
+ " └── Number of tools: 3\n",
167
+ "==================================================\n",
168
+ "Task ID: 6359a0b1-8f7b-499b-9336-840f9ab90688\n",
169
+ "Question: What is the area of the green polygon in the attached file? The numbers in purple represent the lengths of the side they are next to.\n",
170
+ "Level: 2\n",
171
+ "Final Answer: 39\n",
172
+ "Annotator Metadata: \n",
173
+ " β”œβ”€β”€ Steps: \n",
174
+ " β”‚ β”œβ”€β”€ 1. Open the attached file.\n",
175
+ " β”‚ β”œβ”€β”€ 2. Split the shape into five rectangles.\n",
176
+ " β”‚ β”œβ”€β”€ 3. Find the missing side lengths from the side lengths that are given.\n",
177
+ " β”‚ β”œβ”€β”€ 4. Find the area for each rectangle.\n",
178
+ " β”‚ β”œβ”€β”€ 5. Add the areas together to get the area of the entire shape, 39.\n",
179
+ " β”œβ”€β”€ Number of steps: 5\n",
180
+ " β”œβ”€β”€ How long did this take?: 5-10 minutes\n",
181
+ " β”œβ”€β”€ Tools:\n",
182
+ " β”‚ β”œβ”€β”€ 1. Image recognition\n",
183
+ " β”‚ β”œβ”€β”€ 2. OCR\n",
184
+ " β”‚ β”œβ”€β”€ 3. Calculator\n",
185
+ " └── Number of tools: 3\n",
186
+ "==================================================\n",
187
+ "Task ID: b7f857e4-d8aa-4387-af2a-0e844df5b9d8\n",
188
+ "Question: The attached image contains a Python script. Run the Python code against an array of strings, listed below. The output of the Python script will be a URL containing C++ source code. Compile and run this C++ code against the array [35, 12, 8, 99, 21, 5] and return the sum of the third and fifth integers in the sorted list.\n",
189
+ "\n",
190
+ "arr = ['_alg', 'ghi', 'C++', 'jkl', 'tps', '/Q', 'pqr', 'stu', ':', '//', 'rose', 'vwx', 'yz1', '234', 'tta', '567', '890', 'cod', 'e.', 'or', 'g/', 'wiki', '/', 'ing', 'sort', 'abc' , 'or', 'it', 'hms', 'mno' , 'uic', 'ksort', '#', 'ht' ]\n",
191
+ "Level: 2\n",
192
+ "Final Answer: 47\n",
193
+ "Annotator Metadata: \n",
194
+ " β”œβ”€β”€ Steps: \n",
195
+ " β”‚ β”œβ”€β”€ 1. Extract the Python code from the image\n",
196
+ " β”‚ β”œβ”€β”€ 2. Run the code against the provided array. \n",
197
+ " β”‚ β”œβ”€β”€ 3. Navigate to the returned URL (https://web.archive.org/web/20230609112831/https://rosettacode.org/wiki/sorting_algorithms/Quicksort#C++)\n",
198
+ " β”‚ β”œβ”€β”€ 4. Extract the C++ code from the page.\n",
199
+ " β”‚ β”œβ”€β”€ 5. Insert the provided array into the C++ source code:\n",
200
+ " β”‚ β”œβ”€β”€ int main() {\n",
201
+ " β”‚ β”œβ”€β”€ std::vector<int> arr = {35, 12, 8, 99, 21, 5};\n",
202
+ " β”‚ β”œβ”€β”€ quicksort(arr.begin(), arr.end());\n",
203
+ " β”‚ β”œβ”€β”€ for (const auto& num : arr) {\n",
204
+ " β”‚ β”œβ”€β”€ std::cout << num << \" \";\n",
205
+ " β”‚ β”œβ”€β”€ }\n",
206
+ " β”‚ β”œβ”€β”€ std::cout << \"\\n\";\n",
207
+ " β”‚ β”œβ”€β”€ return 0;\n",
208
+ " β”‚ β”œβ”€β”€ }\n",
209
+ " β”‚ β”œβ”€β”€ 6. Compile the edited code.\n",
210
+ " β”‚ β”œβ”€β”€ 7. Run the compiled binary\n",
211
+ " β”œβ”€β”€ Number of steps: 7\n",
212
+ " β”œβ”€β”€ How long did this take?: 45 minutes\n",
213
+ " β”œβ”€β”€ Tools:\n",
214
+ " β”‚ β”œβ”€β”€ 1. File handling\n",
215
+ " β”‚ β”œβ”€β”€ 2. Computer vision or OCR\n",
216
+ " β”‚ β”œβ”€β”€ 3. Web browser\n",
217
+ " β”‚ β”œβ”€β”€ 4. Python\n",
218
+ " β”‚ β”œβ”€β”€ 5. C++ compiler\n",
219
+ " β”‚ β”œβ”€β”€ 6. Calculator \n",
220
+ " └── Number of tools: 6\n",
221
+ "==================================================\n"
222
+ ]
223
+ }
224
+ ],
225
+ "source": [
226
+ "import random\n",
227
+ "\n",
228
+ "random_samples = random.sample(json_QA, 3)\n",
229
+ "for sample in random_samples:\n",
230
+ " print(\"=\" * 50)\n",
231
+ " print(f\"Task ID: {sample['task_id']}\")\n",
232
+ " print(f\"Question: {sample['Question']}\")\n",
233
+ " print(f\"Level: {sample['Level']}\")\n",
234
+ " print(f\"Final Answer: {sample['Final answer']}\")\n",
235
+ " \n",
236
+ " print(f\"Annotator Metadata: \")\n",
237
+ " print(f\" β”œβ”€β”€ Steps: \")\n",
238
+ " for step in sample['Annotator Metadata']['Steps'].split('\\n'):\n",
239
+ " print(f\" β”‚ β”œβ”€β”€ {step}\")\n",
240
+ " print(f\" β”œβ”€β”€ Number of steps: {sample['Annotator Metadata']['Number of steps']}\")\n",
241
+ " print(f\" β”œβ”€β”€ How long did this take?: {sample['Annotator Metadata']['How long did this take?']}\")\n",
242
+ " print(f\" β”œβ”€β”€ Tools:\")\n",
243
+ " for tool in sample['Annotator Metadata']['Tools'].split('\\n'):\n",
244
+ " print(f\" β”‚ β”œβ”€β”€ {tool}\")\n",
245
+ " print(f\" └── Number of tools: {sample['Annotator Metadata']['Number of tools']}\")\n",
246
+ "print(\"=\" * 50)"
247
+ ]
248
+ },
249
+ {
250
+ "cell_type": "markdown",
251
+ "id": "92c785ac",
252
+ "metadata": {},
253
+ "source": [
254
+ "## Create Database \n",
255
+ "\n",
256
+ "- supbase database: https://supabase.com/docs/guides/database/overview "
257
+ ]
258
+ },
259
+ {
260
+ "cell_type": "code",
261
+ "execution_count": null,
262
+ "id": "ac93270f",
263
+ "metadata": {},
264
+ "outputs": [],
265
+ "source": [
266
+ "### Linking to supabase server\n",
267
+ "import os\n",
268
+ "from dotenv import load_dotenv\n",
269
+ "from langchain_huggingface import HuggingFaceEmbeddings\n",
270
+ "from langchain_community.vectorstores import SupabaseVectorStore\n",
271
+ "from supabase.client import Client, create_client\n",
272
+ "\n",
273
+ "\n",
274
+ "load_dotenv()\n",
275
+ "embeddings = HuggingFaceEmbeddings(model_name=\"sentence-transformers/all-mpnet-base-v2\") # dim=768\n",
276
+ "\n",
277
+ "supabase_url = os.environ.get(\"SUPABASE_URL\")\n",
278
+ "supabase_key = os.environ.get(\"SUPABASE_SERVICE_KEY\")\n",
279
+ "supabase: Client = create_client(supabase_url, supabase_key)"
280
+ ]
281
+ },
282
+ {
283
+ "cell_type": "markdown",
284
+ "id": "b745387a",
285
+ "metadata": {},
286
+ "source": [
287
+ "## Embedding and Vector Storage"
288
+ ]
289
+ },
290
+ {
291
+ "cell_type": "code",
292
+ "execution_count": null,
293
+ "id": "fb595623",
294
+ "metadata": {},
295
+ "outputs": [],
296
+ "source": [
297
+ "# Convert question-answer pairs to a list of Document objects for vector database storage\n",
298
+ "from langchain.schema import Document\n",
299
+ "docs = []\n",
300
+ "for sample in json_QA:\n",
301
+ " # Create a formatted text combining question and answer\n",
302
+ " content = f\"Question : {sample['Question']}\\n\\nFinal answer : {sample['Final answer']}\"\n",
303
+ " \n",
304
+ " # Create a document dictionary with content, metadata, and vector embedding\n",
305
+ " doc = {\n",
306
+ " \"content\" : content, # The actual text (question + answer)\n",
307
+ " \"metadata\" : { \n",
308
+ " \"source\" : sample['task_id'] \n",
309
+ " },\n",
310
+ " \"embedding\" : embeddings.embed_query(content), \n",
311
+ " }\n",
312
+ " docs.append(doc)\n",
313
+ "\n",
314
+ "# Insert all documents into Supabase vector database\n",
315
+ "try:\n",
316
+ " response = supabase.table(\"documents\").insert(docs).execute()\n",
317
+ " print(\"Data inserted successfully:\", response)\n",
318
+ "except Exception as e:\n",
319
+ " print(\"Error inserting data into Supabase:\", e)"
320
+ ]
321
+ },
322
+ {
323
+ "cell_type": "markdown",
324
+ "id": "d95564f5",
325
+ "metadata": {},
326
+ "source": [
327
+ "## Vector Store and Create Retriever"
328
+ ]
329
+ },
330
+ {
331
+ "cell_type": "code",
332
+ "execution_count": null,
333
+ "id": "935d2193",
334
+ "metadata": {},
335
+ "outputs": [],
336
+ "source": [
337
+ "# add items to vector database\n",
338
+ "vector_store = SupabaseVectorStore(\n",
339
+ " client=supabase,\n",
340
+ " embedding= embeddings,\n",
341
+ " table_name=\"documents\",\n",
342
+ " query_name=\"match_documents\",\n",
343
+ ")\n",
344
+ "retriever = vector_store.as_retriever()"
345
+ ]
346
+ },
347
+ {
348
+ "cell_type": "markdown",
349
+ "id": "321e3452",
350
+ "metadata": {},
351
+ "source": [
352
+ "## Test Query"
353
+ ]
354
+ },
355
+ {
356
+ "cell_type": "code",
357
+ "execution_count": null,
358
+ "id": "cb568e4f",
359
+ "metadata": {},
360
+ "outputs": [],
361
+ "source": [
362
+ "query = \"On June 6, 2023, an article by Carolyn Collins Petersen was published in Universe Today. This article mentions a team that produced a paper about their observations, linked at the bottom of the article. Find this paper. Under what NASA award number was the work performed by R. G. Arendt supported by?\"\n",
363
+ "\n",
364
+ "# matched_docs = vector_store.similarity_search(query, 2)\n",
365
+ "docs = retriever.invoke(query)\n",
366
+ "docs[0]"
367
+ ]
368
+ },
369
+ {
370
+ "cell_type": "markdown",
371
+ "id": "27e9e2b4",
372
+ "metadata": {},
373
+ "source": [
374
+ "## Tool Usage Frequency in GAIA Benchmark"
375
+ ]
376
+ },
377
+ {
378
+ "cell_type": "code",
379
+ "execution_count": null,
380
+ "id": "5f2de7f1",
381
+ "metadata": {},
382
+ "outputs": [],
383
+ "source": [
384
+ "from collections import Counter, OrderedDict\n",
385
+ "\n",
386
+ "tools = []\n",
387
+ "for sample in json_QA:\n",
388
+ " for tool in sample['Annotator Metadata']['Tools'].split('\\n'):\n",
389
+ " tool = tool[2:].strip().lower()\n",
390
+ " if tool.startswith(\"(\"):\n",
391
+ " tool = tool[11:].strip()\n",
392
+ " tools.append(tool)\n",
393
+ "tools_counter = OrderedDict(Counter(tools))\n",
394
+ "print(\"List of tools used in all samples:\")\n",
395
+ "print(\"Total number of tools used:\", len(tools_counter))\n",
396
+ "for tool, count in tools_counter.items():\n",
397
+ " print(f\" β”œβ”€β”€ {tool}: {count}\")"
398
+ ]
399
+ },
400
+ {
401
+ "cell_type": "markdown",
402
+ "id": "41e913b0",
403
+ "metadata": {},
404
+ "source": [
405
+ "# Graph Implementation\n",
406
+ "\n",
407
+ "Here we build the agent graph with LangGraph to solve those tasks!"
408
+ ]
409
+ },
410
+ {
411
+ "cell_type": "code",
412
+ "execution_count": null,
413
+ "id": "24406a79",
414
+ "metadata": {},
415
+ "outputs": [],
416
+ "source": [
417
+ "system_prompt = \"\"\" \n",
418
+ "You are a helpful assistant tasked with answering questions using a set of tools.\n",
419
+ "If the tool is not available, you can try to find the information online. You can also use your own knowledge to answer the question. \n",
420
+ "You need to provide a step-by-step explanation of how you arrived at the answer.\n",
421
+ "==========================\n",
422
+ "Here is a few examples showing you how to answer the question step by step.\n",
423
+ "\"\"\"\n",
424
+ "for i, samples in enumerate(random_samples):\n",
425
+ " system_prompt += f\"\\nQuestion {i+1}: {samples['Question']}\\nSteps:\\n{samples['Annotator Metadata']['Steps']}\\nTools:\\n{samples['Annotator Metadata']['Tools']}\\nFinal Answer: {samples['Final answer']}\\n\"\n",
426
+ "system_prompt += \"\\n==========================\\n\"\n",
427
+ "system_prompt += \"Now, please answer the following question step by step.\\n\"\n",
428
+ "\n",
429
+ "# save the system_prompt to a file\n",
430
+ "with open('system_prompt.txt', 'w') as f:\n",
431
+ " f.write(system_prompt)"
432
+ ]
433
+ },
434
+ {
435
+ "cell_type": "code",
436
+ "execution_count": null,
437
+ "id": "15ab2939",
438
+ "metadata": {},
439
+ "outputs": [],
440
+ "source": [
441
+ "# load the system prompt from the file\n",
442
+ "with open('system_prompt.txt', 'r') as f:\n",
443
+ " system_prompt = f.read()\n",
444
+ "print(system_prompt) "
445
+ ]
446
+ },
447
+ {
448
+ "cell_type": "code",
449
+ "execution_count": null,
450
+ "id": "5e33194b",
451
+ "metadata": {},
452
+ "outputs": [],
453
+ "source": [
454
+ "import dotenv\n",
455
+ "from langgraph.graph import MessagesState, START, StateGraph\n",
456
+ "from langgraph.prebuilt import tools_condition\n",
457
+ "from langgraph.prebuilt import ToolNode\n",
458
+ "from langchain_google_genai import ChatGoogleGenerativeAI\n",
459
+ "from langchain_core.messages import HumanMessage, SystemMessage\n",
460
+ "from tools.searchtools import question_retrieve_tool, wiki_search, web_search, arvix_search\n",
461
+ "from tools.mathtools import multiply, add, subtract, divide, modulus\n",
462
+ "\n",
463
+ "# Define the retriever from supabase\n",
464
+ "load_dotenv()\n",
465
+ "\n",
466
+ "tools = [\n",
467
+ " multiply,\n",
468
+ " add,\n",
469
+ " subtract,\n",
470
+ " divide,\n",
471
+ " modulus,\n",
472
+ " wiki_search,\n",
473
+ " web_search,\n",
474
+ " arvix_search,\n",
475
+ " question_retrieve_tool\n",
476
+ "]\n",
477
+ "\n",
478
+ "llm = ChatGoogleGenerativeAI(model=\"gemini-2.0-flash\")\n",
479
+ "llm_with_tools = llm.bind_tools(tools)"
480
+ ]
481
+ },
482
+ {
483
+ "cell_type": "code",
484
+ "execution_count": null,
485
+ "id": "2e8901db",
486
+ "metadata": {},
487
+ "outputs": [],
488
+ "source": [
489
+ "# load the system prompt from the file\n",
490
+ "with open('system_prompt.txt', 'r') as f:\n",
491
+ " system_prompt = f.read()\n",
492
+ "\n",
493
+ "\n",
494
+ "# System message\n",
495
+ "sys_msg = SystemMessage(content=system_prompt)\n",
496
+ "\n",
497
+ "# Node\n",
498
+ "def assistant(state: MessagesState):\n",
499
+ " \"\"\"Assistant node\"\"\"\n",
500
+ " return {\"messages\": [llm_with_tools.invoke([sys_msg] + state[\"messages\"])]}\n",
501
+ "\n",
502
+ "# Build graph\n",
503
+ "builder = StateGraph(MessagesState)\n",
504
+ "builder.add_node(\"assistant\", assistant)\n",
505
+ "builder.add_node(\"tools\", ToolNode(tools))\n",
506
+ "builder.add_edge(START, \"assistant\")\n",
507
+ "builder.add_conditional_edges(\n",
508
+ " \"assistant\",\n",
509
+ " # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools\n",
510
+ " # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END\n",
511
+ " tools_condition,\n",
512
+ ")\n",
513
+ "builder.add_edge(\"tools\", \"assistant\")\n",
514
+ "\n",
515
+ "# Compile graph\n",
516
+ "graph = builder.compile()"
517
+ ]
518
+ },
519
+ {
520
+ "cell_type": "code",
521
+ "execution_count": null,
522
+ "id": "74ba4cd8",
523
+ "metadata": {},
524
+ "outputs": [],
525
+ "source": [
526
+ "from IPython.display import Image, display\n",
527
+ "\n",
528
+ "display(Image(graph.get_graph(xray=True).draw_mermaid_png()))"
529
+ ]
530
+ },
531
+ {
532
+ "cell_type": "code",
533
+ "execution_count": null,
534
+ "id": "af40c038",
535
+ "metadata": {},
536
+ "outputs": [],
537
+ "source": [
538
+ "question = \"On June 6, 2023, an article by Carolyn Collins Petersen was published in Universe Today. This article mentions a team that produced a paper about their observations, linked at the bottom of the article. Find this paper. Under what NASA award number was the work performed by R. G. Arendt supported by?\"\n",
539
+ "messages = [HumanMessage(content=question)]\n",
540
+ "messages = graph.invoke({\"messages\": messages})"
541
+ ]
542
+ },
543
+ {
544
+ "cell_type": "code",
545
+ "execution_count": null,
546
+ "id": "133aba59",
547
+ "metadata": {},
548
+ "outputs": [],
549
+ "source": [
550
+ "for m in messages['messages']:\n",
551
+ " m.pretty_print()"
552
+ ]
553
+ },
554
+ {
555
+ "cell_type": "code",
556
+ "execution_count": null,
557
+ "id": "f6738871",
558
+ "metadata": {},
559
+ "outputs": [],
560
+ "source": []
561
+ }
562
+ ],
563
+ "metadata": {
564
+ "kernelspec": {
565
+ "display_name": "Python 3",
566
+ "language": "python",
567
+ "name": "python3"
568
+ },
569
+ "language_info": {
570
+ "codemirror_mode": {
571
+ "name": "ipython",
572
+ "version": 3
573
+ },
574
+ "file_extension": ".py",
575
+ "mimetype": "text/x-python",
576
+ "name": "python",
577
+ "nbconvert_exporter": "python",
578
+ "pygments_lexer": "ipython3",
579
+ "version": "3.12.6"
580
+ }
581
+ },
582
+ "nbformat": 4,
583
+ "nbformat_minor": 5
584
+ }
tools/__init__.py ADDED
File without changes