Update agent file

README.md

@@ -1,15 +1,65 @@
----
-title: Template Final Assignment
-emoji: 🕵🏻‍♂️
-colorFrom: indigo
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.25.2
-app_file: app.py
-pinned: false
-hf_oauth: true
-# optional, default duration is 8 hours/480 minutes. Max duration is 30 days/43200 minutes.
-hf_oauth_expiration_minutes: 480
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+
+## 1. Setup Supabase for Vector Store, following [run this SQL Editor](https://supabase.com/dashboard/project/upkfycxsetvrlochkrti/sql/c412918f-ec41-4030-9261-80bc415e247e) 
+
+- enable `vector` extension;
+- create `documents` table;
+- create `match_documents_langchain` function.
+
+```
+-- Enable the pgvector extension to work with embedding vectors
+create extension vector;
+
+-- Create a table to store your documents
+create table documents (
+  id bigserial primary key,
+  content text, -- corresponds to Document.pageContent
+  metadata jsonb, -- corresponds to Document.metadata
+  embedding vector -- 1536 works for OpenAI embeddings, change if needed
+);
+
+-- Create a function to search for documents
+create function match_documents_langchain (
+  query_embedding vector,
+  match_count int default null,
+  filter jsonb DEFAULT '{}'
+) returns table (
+  id bigint,
+  content text,
+  metadata jsonb,
+  similarity float
+)
+language plpgsql
+as $$
+#variable_conflict use_column
+begin
+  return query
+  select
+    id,
+    content,
+    metadata,
+    1 - (documents.embedding <=> query_embedding) as similarity
+  from documents
+  where metadata @> filter
+  order by documents.embedding <=> query_embedding
+  limit match_count;
+end;
+$$;
+```
+
+## 2. Setup Supabase API Key
+
+```shell
+export SUPABASE_URL=https://upkfycxsetvrlochkrti.supabase.co
+export SUPABASE_SERVICE_KEY=
+```
+
+## 3. Setup Google Gemini API Key, or Groq Cloud
+
+```shell
+# Google Gemini
+export GOOGLE_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxx_xxxxxxxx
+# Groq Cloud
+export GROQ_API_KEY=gsk_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+```
+

agent.py

@@ -18,9 +18,11 @@ from supabase.client import Client, create_client
 
 load_dotenv()
 
+
 @tool
 def multiply(a: int, b: int) -> int:
     """Multiply two numbers.
+
     Args:
         a: first int
         b: second int
@@ -30,7 +32,7 @@ def multiply(a: int, b: int) -> int:
 @tool
 def add(a: int, b: int) -> int:
     """Add two numbers.
-    
+
     Args:
         a: first int
         b: second int
@@ -40,7 +42,7 @@ def add(a: int, b: int) -> int:
 @tool
 def subtract(a: int, b: int) -> int:
     """Subtract two numbers.
-    
+
     Args:
         a: first int
         b: second int
@@ -50,7 +52,7 @@ def subtract(a: int, b: int) -> int:
 @tool
 def divide(a: int, b: int) -> int:
     """Divide two numbers.
-    
+
     Args:
         a: first int
         b: second int
@@ -62,7 +64,7 @@ def divide(a: int, b: int) -> int:
 @tool
 def modulus(a: int, b: int) -> int:
     """Get the modulus of two numbers.
-    
+
     Args:
         a: first int
         b: second int
@@ -72,7 +74,7 @@ def modulus(a: int, b: int) -> int:
 @tool
 def wiki_search(query: str) -> str:
     """Search Wikipedia for a query and return maximum 2 results.
-    
+
     Args:
         query: The search query."""
     search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
@@ -86,7 +88,7 @@ def wiki_search(query: str) -> str:
 @tool
 def web_search(query: str) -> str:
     """Search Tavily for a query and return maximum 3 results.
-    
+
     Args:
         query: The search query."""
     search_docs = TavilySearchResults(max_results=3).invoke(query=query)
@@ -100,7 +102,7 @@ def web_search(query: str) -> str:
 @tool
 def arvix_search(query: str) -> str:
     """Search Arxiv for a query and return maximum 3 result.
-    
+
     Args:
         query: The search query."""
     search_docs = ArxivLoader(query=query, load_max_docs=3).load()
@@ -111,31 +113,36 @@ def arvix_search(query: str) -> str:
         ])
     return {"arvix_results": formatted_search_docs}
 
+@tool
+def similar_question_search(question: str) -> str:
+    """Search the vector database for similar questions and return the first results.
 
+    Args:
+        question: the question human provided."""
+    matched_docs = vector_store.similarity_search(query, 3)
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+            for doc in matched_docs
+        ])
+    return {"similar_questions": formatted_search_docs}
 
-# load the system prompt from the file
-with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    system_prompt = f.read()
-
-# System message
-sys_msg = SystemMessage(content=system_prompt)
 
-# build a retriever
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") #  dim=768
-supabase: Client = create_client(
-    os.environ.get("SUPABASE_URL"),
-    os.environ.get("SUPABASE_SERVICE_KEY"))
+supabase_url = os.environ.get("SUPABASE_URL")
+supabase_key = os.environ.get("SUPABASE_SERVICE_KEY")
+supabase: Client = create_client(supabase_url, supabase_key)
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 vector_store = SupabaseVectorStore(
     client=supabase,
     embedding= embeddings,
     table_name="documents",
     query_name="match_documents_langchain",
 )
-create_retriever_tool = create_retriever_tool(
-    retriever=vector_store.as_retriever(),
-    name="Question Search",
-    description="A tool to retrieve similar questions from a vector store.",
-)
+# question_retrieve_tool = create_retriever_tool(
+#     vector_store.as_retriever(),
+#     "Question Retriever",
+#     "Find similar questions in the vector database for the given question.",
+# )
 
 tools = [
     multiply,
@@ -145,36 +152,26 @@ tools = [
     modulus,
     wiki_search,
     web_search,
-    arvix_search,
+    arvix_search
 ]
 
-# Build graph function
-def build_graph(provider: str = "groq"):
-    """Build the graph"""
-    # Load environment variables from .env file
-    if provider == "google":
-        # Google Gemini
-        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-    elif provider == "groq":
-        # Groq https://console.groq.com/docs/models
-        llm = ChatGroq(model="qwen-qwq-32b", temperature=0) # optional : qwen-qwq-32b gemma2-9b-it
-    elif provider == "huggingface":
-        # TODO: Add huggingface endpoint
-        llm = ChatHuggingFace(
-            llm=HuggingFaceEndpoint(
-                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
-                temperature=0,
-            ),
-        )
-    else:
-        raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
-    # Bind tools to LLM
+
+
+# load the system prompt from the file
+with open('system_prompt.txt', 'r') as f:
+    system_prompt = f.read()
+# System message
+sys_msg = SystemMessage(content=system_prompt)
+
+
+def build_graph():
+    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-preview-05-20")
     llm_with_tools = llm.bind_tools(tools)
 
     # Node
     def assistant(state: MessagesState):
         """Assistant node"""
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+        return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
 
     def retriever(state: MessagesState):
         """Retriever node"""
@@ -184,6 +181,7 @@ def build_graph(provider: str = "groq"):
         )
         return {"messages": [sys_msg] + state["messages"] + [example_msg]}
 
+    # Build graph
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
@@ -192,20 +190,27 @@ def build_graph(provider: str = "groq"):
     builder.add_edge("retriever", "assistant")
     builder.add_conditional_edges(
         "assistant",
+        # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
+        # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
         tools_condition,
     )
     builder.add_edge("tools", "assistant")
 
     # Compile graph
-    return builder.compile()
+    graph = builder.compile()
+
+    # from IPython.display import Image, display
+    # display(Image(graph.get_graph(xray=True).draw_mermaid_png()))
+
+    return graph
+
 
-# test
 if __name__ == "__main__":
     question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
     # Build the graph
-    graph = build_graph(provider="groq")
+    graph = build_graph()
     # Run the graph
     messages = [HumanMessage(content=question)]
     messages = graph.invoke({"messages": messages})
     for m in messages["messages"]:
-        m.pretty_print()
+        m.pretty_print()

app.py

@@ -21,7 +21,7 @@ class BasicAgent:
     """A langgraph agent."""
     def __init__(self):
         print("BasicAgent initialized.")
-        self.graph = build_graph(provider="google")
+        self.graph = build_graph()
 
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")

code_interpreter.py

@@ -1,281 +0,0 @@
-import os
-import io
-import sys
-import uuid
-import base64
-import traceback
-import contextlib
-import tempfile
-import subprocess
-import sqlite3
-from typing import Dict, List, Any, Optional, Union
-import numpy as np
-import pandas as pd
-import matplotlib.pyplot as plt
-from PIL import Image
-
-class CodeInterpreter:
-    def __init__(self, allowed_modules=None, max_execution_time=30, working_directory=None):
-        """Initialize the code interpreter with safety measures."""
-        self.allowed_modules = allowed_modules or [
-            "numpy", "pandas", "matplotlib", "scipy", "sklearn", 
-            "math", "random", "statistics", "datetime", "collections",
-            "itertools", "functools", "operator", "re", "json",
-            "sympy", "networkx", "nltk", "PIL", "pytesseract", 
-            "cmath", "uuid", "tempfile", "requests", "urllib"
-        ]
-        self.max_execution_time = max_execution_time
-        self.working_directory = working_directory or os.path.join(os.getcwd()) 
-        if not os.path.exists(self.working_directory):
-            os.makedirs(self.working_directory)
-        
-        self.globals = {
-            "__builtins__": __builtins__,
-            "np": np,
-            "pd": pd,
-            "plt": plt,
-            "Image": Image,
-        }
-        self.temp_sqlite_db = os.path.join(tempfile.gettempdir(), "code_exec.db")
-
-    def execute_code(self, code: str, language: str = "python") -> Dict[str, Any]:
-        """Execute the provided code in the selected programming language."""
-        language = language.lower()
-        execution_id = str(uuid.uuid4())
-        
-        result = {
-            "execution_id": execution_id,
-            "status": "error",
-            "stdout": "",
-            "stderr": "",
-            "result": None,
-            "plots": [],
-            "dataframes": []
-        }
-        
-        try:
-            if language == "python":
-                return self._execute_python(code, execution_id)
-            elif language == "bash":
-                return self._execute_bash(code, execution_id)
-            elif language == "sql":
-                return self._execute_sql(code, execution_id)
-            elif language == "c":
-                return self._execute_c(code, execution_id)
-            elif language == "java":
-                return self._execute_java(code, execution_id)
-            else:
-                result["stderr"] = f"Unsupported language: {language}"
-        except Exception as e:
-            result["stderr"] = str(e)
-        
-        return result
-
-    def _execute_python(self, code: str, execution_id: str) -> dict:
-        output_buffer = io.StringIO()
-        error_buffer = io.StringIO()
-        result = {
-            "execution_id": execution_id,
-            "status": "error",
-            "stdout": "",
-            "stderr": "",
-            "result": None,
-            "plots": [],
-            "dataframes": []
-        }
-        
-        try:
-            exec_dir = os.path.join(self.working_directory, execution_id)
-            os.makedirs(exec_dir, exist_ok=True)
-            plt.switch_backend('Agg')
-            
-            with contextlib.redirect_stdout(output_buffer), contextlib.redirect_stderr(error_buffer):
-                exec_result = exec(code, self.globals)
-
-                if plt.get_fignums():
-                    for i, fig_num in enumerate(plt.get_fignums()):
-                        fig = plt.figure(fig_num)
-                        img_path = os.path.join(exec_dir, f"plot_{i}.png")
-                        fig.savefig(img_path)
-                        with open(img_path, "rb") as img_file:
-                            img_data = base64.b64encode(img_file.read()).decode('utf-8')
-                            result["plots"].append({
-                                "figure_number": fig_num,
-                                "data": img_data
-                            })
-
-                for var_name, var_value in self.globals.items():
-                    if isinstance(var_value, pd.DataFrame) and len(var_value) > 0:
-                        result["dataframes"].append({
-                            "name": var_name,
-                            "head": var_value.head().to_dict(),
-                            "shape": var_value.shape,
-                            "dtypes": str(var_value.dtypes)
-                        })
-                
-            result["status"] = "success"
-            result["stdout"] = output_buffer.getvalue()
-            result["result"] = exec_result
-            
-        except Exception as e:
-            result["status"] = "error"
-            result["stderr"] = f"{error_buffer.getvalue()}\n{traceback.format_exc()}"
-        
-        return result
-
-    def _execute_bash(self, code: str, execution_id: str) -> dict:
-        try:
-            completed = subprocess.run(
-                code, shell=True, capture_output=True, text=True, timeout=self.max_execution_time
-            )
-            return {
-                "execution_id": execution_id,
-                "status": "success" if completed.returncode == 0 else "error",
-                "stdout": completed.stdout,
-                "stderr": completed.stderr,
-                "result": None,
-                "plots": [],
-                "dataframes": []
-            }
-        except subprocess.TimeoutExpired:
-            return {
-                "execution_id": execution_id,
-                "status": "error",
-                "stdout": "",
-                "stderr": "Execution timed out.",
-                "result": None,
-                "plots": [],
-                "dataframes": []
-            }
-
-    def _execute_sql(self, code: str, execution_id: str) -> dict:
-        result = {
-            "execution_id": execution_id,
-            "status": "error",
-            "stdout": "",
-            "stderr": "",
-            "result": None,
-            "plots": [],
-            "dataframes": []
-        }
-        try:
-            conn = sqlite3.connect(self.temp_sqlite_db)
-            cur = conn.cursor()
-            cur.execute(code)
-            if code.strip().lower().startswith("select"):
-                columns = [description[0] for description in cur.description]
-                rows = cur.fetchall()
-                df = pd.DataFrame(rows, columns=columns)
-                result["dataframes"].append({
-                    "name": "query_result",
-                    "head": df.head().to_dict(),
-                    "shape": df.shape,
-                    "dtypes": str(df.dtypes)
-                })
-            else:
-                conn.commit()
-
-            result["status"] = "success"
-            result["stdout"] = "Query executed successfully."
-
-        except Exception as e:
-            result["stderr"] = str(e)
-        finally:
-            conn.close()
-
-        return result
-
-    def _execute_c(self, code: str, execution_id: str) -> dict:
-        temp_dir = tempfile.mkdtemp()
-        source_path = os.path.join(temp_dir, "program.c")
-        binary_path = os.path.join(temp_dir, "program")
-
-        try:
-            with open(source_path, "w") as f:
-                f.write(code)
-
-            compile_proc = subprocess.run(
-                ["gcc", source_path, "-o", binary_path],
-                capture_output=True, text=True, timeout=self.max_execution_time
-            )
-            if compile_proc.returncode != 0:
-                return {
-                    "execution_id": execution_id,
-                    "status": "error",
-                    "stdout": compile_proc.stdout,
-                    "stderr": compile_proc.stderr,
-                    "result": None,
-                    "plots": [],
-                    "dataframes": []
-                }
-
-            run_proc = subprocess.run(
-                [binary_path],
-                capture_output=True, text=True, timeout=self.max_execution_time
-            )
-            return {
-                "execution_id": execution_id,
-                "status": "success" if run_proc.returncode == 0 else "error",
-                "stdout": run_proc.stdout,
-                "stderr": run_proc.stderr,
-                "result": None,
-                "plots": [],
-                "dataframes": []
-            }
-        except Exception as e:
-            return {
-                "execution_id": execution_id,
-                "status": "error",
-                "stdout": "",
-                "stderr": str(e),
-                "result": None,
-                "plots": [],
-                "dataframes": []
-            }
-
-    def _execute_java(self, code: str, execution_id: str) -> dict:
-        temp_dir = tempfile.mkdtemp()
-        source_path = os.path.join(temp_dir, "Main.java")
-
-        try:
-            with open(source_path, "w") as f:
-                f.write(code)
-
-            compile_proc = subprocess.run(
-                ["javac", source_path],
-                capture_output=True, text=True, timeout=self.max_execution_time
-            )
-            if compile_proc.returncode != 0:
-                return {
-                    "execution_id": execution_id,
-                    "status": "error",
-                    "stdout": compile_proc.stdout,
-                    "stderr": compile_proc.stderr,
-                    "result": None,
-                    "plots": [],
-                    "dataframes": []
-                }
-
-            run_proc = subprocess.run(
-                ["java", "-cp", temp_dir, "Main"],
-                capture_output=True, text=True, timeout=self.max_execution_time
-            )
-            return {
-                "execution_id": execution_id,
-                "status": "success" if run_proc.returncode == 0 else "error",
-                "stdout": run_proc.stdout,
-                "stderr": run_proc.stderr,
-                "result": None,
-                "plots": [],
-                "dataframes": []
-            }
-        except Exception as e:
-            return {
-                "execution_id": execution_id,
-                "status": "error",
-                "stdout": "",
-                "stderr": str(e),
-                "result": None,
-                "plots": [],
-                "dataframes": []
-            }

explore_metadata.ipynb

@@ -1,336 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "id": "a600d7fc",
-   "metadata": {
-    "jupyter": {
-     "is_executing": true
-    }
-   },
-   "source": [
-    "import json \n",
-    "with open('metadata.jsonl', 'r') as f: \n",
-    "    json_list = list(f)\n",
-    "\n",
-    "json_QA = []\n",
-    "for json_str in json_list: \n",
-    "    json_data = json.loads(json_str)\n",
-    "    json_QA.append(json_data)"
-   ],
-   "outputs": [],
-   "execution_count": null
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "id": "fa5d8eb8",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "==================================================\n",
-      "Task ID: d1af70ea-a9a4-421a-b9cc-94b5e02f1788\n",
-      "Question: As of the 2020 census, what was the population difference between the largest county seat and smallest county seat, by land area of the county seat, in Washington state? For population figures, please use the official data from data.census.gov. Please report the integer difference.\n",
-      "Level: 2\n",
-      "Final Answer: 736455\n",
-      "Annotator Metadata: \n",
-      "  ├── Steps: \n",
-      "  │      ├── Step 1: Using a web browser, access a search engine and conduct a search, \"Washington cities by area\"\n",
-      "  │      ├── Step 2: Navigate to the second search result, https://en.wikipedia.org/wiki/List_of_municipalities_in_Washington\n",
-      "  │      ├── Step 3: Evaluate the page contents, finding the largest and smallest county seats by land area, Seattle and Cathlamet\n",
-      "  │      ├── Step 4: Using a web browser, navigate to https://data.census.gov/\n",
-      "  │      ├── Step 5: Using the website's search area, conduct a search, Seattle, Washington\n",
-      "  │      ├── Step 6: Record the reported 2020 Decennial Census population of Seattle, Washington, 737,015\n",
-      "  │      ├── Step 7: Using the website's search area, conduct a search, Cathlamet, Washington\n",
-      "  │      ├── Step 8: Record the reported 2020 Decennial Census population of Cathlamet, Washington, 560\n",
-      "  │      ├── Step 9: Using a calculator, find the difference in populations,\n",
-      "  │      ├── \n",
-      "  │      ├── 737,015 - 560\n",
-      "  │      ├── 736,455\n",
-      "  │      ├── Step 10: Report the correct answer to my user in the requested format, \"736,455\"\n",
-      "  ├── Number of steps: 10\n",
-      "  ├── How long did this take?: 5 minutes\n",
-      "  ├── Tools:\n",
-      "  │      ├── 1. A web browser\n",
-      "  │      ├── 2. A search engine\n",
-      "  │      ├── 3. A calculator\n",
-      "  └── Number of tools: 3\n",
-      "==================================================\n"
-     ]
-    }
-   ],
-   "source": [
-    "import random\n",
-    "random_samples = random.sample(json_QA, 1)\n",
-    "for sample in random_samples:\n",
-    "    print(\"=\" * 50)\n",
-    "    print(f\"Task ID: {sample['task_id']}\")\n",
-    "    print(f\"Question: {sample['Question']}\")\n",
-    "    print(f\"Level: {sample['Level']}\")\n",
-    "    print(f\"Final Answer: {sample['Final answer']}\")\n",
-    "    print(f\"Annotator Metadata: \")\n",
-    "    print(f\"  ├── Steps: \")\n",
-    "    for step in sample['Annotator Metadata']['Steps'].split('\\n'):\n",
-    "        print(f\"  │      ├── {step}\")\n",
-    "    print(f\"  ├── Number of steps: {sample['Annotator Metadata']['Number of steps']}\")\n",
-    "    print(f\"  ├── How long did this take?: {sample['Annotator Metadata']['How long did this take?']}\")\n",
-    "    print(f\"  ├── Tools:\")\n",
-    "    for tool in sample['Annotator Metadata']['Tools'].split('\\n'):\n",
-    "        print(f\"  │      ├── {tool}\")\n",
-    "    print(f\"  └── Number of tools: {sample['Annotator Metadata']['Number of tools']}\")\n",
-    "print(\"=\" * 50)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "id": "05076516",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "from dotenv import load_dotenv\n",
-    "from langchain_huggingface import HuggingFaceEmbeddings\n",
-    "from langchain_community.vectorstores import SupabaseVectorStore\n",
-    "from supabase.client import Client, create_client\n",
-    "\n",
-    "\n",
-    "load_dotenv()\n",
-    "embeddings = HuggingFaceEmbeddings(model_name=\"sentence-transformers/all-mpnet-base-v2\") #  dim=768\n",
-    "\n",
-    "supabase_url = os.environ.get(\"SUPABASE_URL\")\n",
-    "supabase_key = os.environ.get(\"SUPABASE_SERVICE_ROLE_KEY\")\n",
-    "supabase: Client = create_client(supabase_url, supabase_key)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 20,
-   "id": "aa1402e3",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from langchain.schema import Document\n",
-    "docs = []\n",
-    "cnt = 0 \n",
-    "for sample in json_QA:\n",
-    "    content = f\"Question : {sample['Question']}\\n\\nFinal answer : {sample['Final answer']}\"\n",
-    "    doc = {\n",
-    "        \"id\" : cnt,\n",
-    "        \"content\" : content,\n",
-    "        \"metadata\" : {\n",
-    "            \"source\" : sample['task_id']\n",
-    "        },\n",
-    "        \"embedding\" : embeddings.embed_query(content),\n",
-    "    }\n",
-    "    docs.append(doc)\n",
-    "    cnt += 1\n",
-    "\n",
-    "# upload the documents to the vector database\n",
-    "try:\n",
-    "    response = (\n",
-    "        supabase.table(\"documents2\")\n",
-    "        .insert(docs)\n",
-    "        .execute()\n",
-    "    )\n",
-    "except Exception as exception:\n",
-    "    print(\"Error inserting data into Supabase:\", exception)\n",
-    "\n",
-    "# # Save the documents (a list of dict) into a csv file, and manually upload it to Supabase\n",
-    "# import pandas as pd\n",
-    "# df = pd.DataFrame(docs)\n",
-    "# df.to_csv('supabase_docs.csv',index=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 41,
-   "id": "9aa7eb5e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# add items to vector database\n",
-    "vector_store = SupabaseVectorStore(\n",
-    "    client=supabase,\n",
-    "    embedding= embeddings,\n",
-    "    table_name=\"documents2\",\n",
-    "    query_name=\"match_documents_2\",\n",
-    ")\n",
-    "retriever = vector_store.as_retriever()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 42,
-   "id": "9eecafd1",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "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",
-    "# matched_docs = vector_store.similarity_search(query, k=2)\n",
-    "docs = retriever.invoke(query)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 43,
-   "id": "ff917840",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Document(metadata={'source': '840bfca7-4f7b-481a-8794-c560c340185d'}, page_content='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\\nFinal answer : 80GSFC21M0002')"
-      ]
-     },
-     "execution_count": 43,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "docs[0]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 44,
-   "id": "01c8f337",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "List of tools used in all samples:\n",
-      "Total number of tools used: 83\n",
-      "  ├── web browser: 107\n",
-      "  ├── image recognition tools (to identify and parse a figure with three axes): 1\n",
-      "  ├── search engine: 101\n",
-      "  ├── calculator: 34\n",
-      "  ├── unlambda compiler (optional): 1\n",
-      "  ├── a web browser.: 2\n",
-      "  ├── a search engine.: 2\n",
-      "  ├── a calculator.: 1\n",
-      "  ├── microsoft excel: 5\n",
-      "  ├── google search: 1\n",
-      "  ├── ne: 9\n",
-      "  ├── pdf access: 7\n",
-      "  ├── file handling: 2\n",
-      "  ├── python: 3\n",
-      "  ├── image recognition tools: 12\n",
-      "  ├── jsonld file access: 1\n",
-      "  ├── video parsing: 1\n",
-      "  ├── python compiler: 1\n",
-      "  ├── video recognition tools: 3\n",
-      "  ├── pdf viewer: 7\n",
-      "  ├── microsoft excel / google sheets: 3\n",
-      "  ├── word document access: 1\n",
-      "  ├── tool to extract text from images: 1\n",
-      "  ├── a word reversal tool / script: 1\n",
-      "  ├── counter: 1\n",
-      "  ├── excel: 3\n",
-      "  ├── image recognition: 5\n",
-      "  ├── color recognition: 3\n",
-      "  ├── excel file access: 3\n",
-      "  ├── xml file access: 1\n",
-      "  ├── access to the internet archive, web.archive.org: 1\n",
-      "  ├── text processing/diff tool: 1\n",
-      "  ├── gif parsing tools: 1\n",
-      "  ├── a web browser: 7\n",
-      "  ├── a search engine: 7\n",
-      "  ├── a speech-to-text tool: 2\n",
-      "  ├── code/data analysis tools: 1\n",
-      "  ├── audio capability: 2\n",
-      "  ├── pdf reader: 1\n",
-      "  ├── markdown: 1\n",
-      "  ├── a calculator: 5\n",
-      "  ├── access to wikipedia: 3\n",
-      "  ├── image recognition/ocr: 3\n",
-      "  ├── google translate access: 1\n",
-      "  ├── ocr: 4\n",
-      "  ├── bass note data: 1\n",
-      "  ├── text editor: 1\n",
-      "  ├── xlsx file access: 1\n",
-      "  ├── powerpoint viewer: 1\n",
-      "  ├── csv file access: 1\n",
-      "  ├── calculator (or use excel): 1\n",
-      "  ├── computer algebra system: 1\n",
-      "  ├── video processing software: 1\n",
-      "  ├── audio processing software: 1\n",
-      "  ├── computer vision: 1\n",
-      "  ├── google maps: 1\n",
-      "  ├── access to excel files: 1\n",
-      "  ├── calculator (or ability to count): 1\n",
-      "  ├── a file interface: 3\n",
-      "  ├── a python ide: 1\n",
-      "  ├── spreadsheet editor: 1\n",
-      "  ├── tools required: 1\n",
-      "  ├── b browser: 1\n",
-      "  ├── image recognition and processing tools: 1\n",
-      "  ├── computer vision or ocr: 1\n",
-      "  ├── c++ compiler: 1\n",
-      "  ├── access to google maps: 1\n",
-      "  ├── youtube player: 1\n",
-      "  ├── natural language processor: 1\n",
-      "  ├── graph interaction tools: 1\n",
-      "  ├── bablyonian cuniform -> arabic legend: 1\n",
-      "  ├── access to youtube: 1\n",
-      "  ├── image search tools: 1\n",
-      "  ├── calculator or counting function: 1\n",
-      "  ├── a speech-to-text audio processing tool: 1\n",
-      "  ├── access to academic journal websites: 1\n",
-      "  ├── pdf reader/extracter: 1\n",
-      "  ├── rubik's cube model: 1\n",
-      "  ├── wikipedia: 1\n",
-      "  ├── video capability: 1\n",
-      "  ├── image processing tools: 1\n",
-      "  ├── age recognition software: 1\n",
-      "  ├── youtube: 1\n"
-     ]
-    }
-   ],
-   "source": [
-    "# list of the tools used in all the samples\n",
-    "from collections import Counter, OrderedDict\n",
-    "\n",
-    "tools = []\n",
-    "for sample in json_QA:\n",
-    "    for tool in sample['Annotator Metadata']['Tools'].split('\\n'):\n",
-    "        tool = tool[2:].strip().lower()\n",
-    "        if tool.startswith(\"(\"):\n",
-    "            tool = tool[11:].strip()\n",
-    "        tools.append(tool)\n",
-    "tools_counter = OrderedDict(Counter(tools))\n",
-    "print(\"List of tools used in all samples:\")\n",
-    "print(\"Total number of tools used:\", len(tools_counter))\n",
-    "for tool, count in tools_counter.items():\n",
-    "    print(f\"  ├── {tool}: {count}\")"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "env",
-   "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.11.9"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}

image_processing.py

@@ -1,26 +0,0 @@
-import os
-import io
-import base64
-import uuid
-from PIL import Image
-
-# Helper functions for image processing
-def encode_image(image_path: str) -> str:
-    """Convert an image file to base64 string."""
-    with open(image_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode("utf-8")
-
-
-def decode_image(base64_string: str) -> Image.Image:
-    """Convert a base64 string to a PIL Image."""
-    image_data = base64.b64decode(base64_string)
-    return Image.open(io.BytesIO(image_data))
-
-
-def save_image(image: Image.Image, directory: str = "image_outputs") -> str:
-    """Save a PIL Image to disk and return the path."""
-    os.makedirs(directory, exist_ok=True)
-    image_id = str(uuid.uuid4())
-    image_path = os.path.join(directory, f"{image_id}.png")
-    image.save(image_path)
-    return image_path

system_prompt.txt

@@ -14,4 +14,4 @@ Examples:
 - FINAL ANSWER: Paris
 - FINAL ANSWER: 128
 
-If you do not follow this format exactly, your response will be considered incorrect.
+If you do not follow this format exactly, your response will be considered incorrect.

test.ipynb

@@ -0,0 +1,459 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "id": "initial_id",
+   "metadata": {
+    "collapsed": true
+   },
+   "source": [
+    "# Load metadata.jsonl\n",
+    "import json\n",
+    "# Load the metadata.jsonl file\n",
+    "with open('metadata.jsonl', 'r') as jsonl_file:\n",
+    "    json_list = list(jsonl_file)\n",
+    "\n",
+    "json_QA = []\n",
+    "for json_str in json_list:\n",
+    "    json_data = json.loads(json_str)\n",
+    "    json_QA.append(json_data)"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "# randomly select 3 samples\n",
+    "# {\"task_id\": \"c61d22de-5f6c-4958-a7f6-5e9707bd3466\", \"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?\", \"Level\": 2, \"Final answer\": \"egalitarian\", \"file_name\": \"\", \"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.\", \"Number of steps\": \"12\", \"How long did this take?\": \"8 minutes\", \"Tools\": \"1. Web browser\\n2. Image recognition tools (to identify and parse a figure with three axes)\", \"Number of tools\": \"2\"}}\n",
+    "\n",
+    "import random\n",
+    "# random.seed(42)\n",
+    "random_samples = random.sample(json_QA, 1)\n",
+    "for sample in random_samples:\n",
+    "    print(\"=\" * 50)\n",
+    "    print(f\"Task ID: {sample['task_id']}\")\n",
+    "    print(f\"Question: {sample['Question']}\")\n",
+    "    print(f\"Level: {sample['Level']}\")\n",
+    "    print(f\"Final Answer: {sample['Final answer']}\")\n",
+    "    print(f\"Annotator Metadata: \")\n",
+    "    print(f\"  ├── Steps: \")\n",
+    "    for step in sample['Annotator Metadata']['Steps'].split('\\n'):\n",
+    "        print(f\"  │      ├── {step}\")\n",
+    "    print(f\"  ├── Number of steps: {sample['Annotator Metadata']['Number of steps']}\")\n",
+    "    print(f\"  ├── How long did this take?: {sample['Annotator Metadata']['How long did this take?']}\")\n",
+    "    print(f\"  ├── Tools:\")\n",
+    "    for tool in sample['Annotator Metadata']['Tools'].split('\\n'):\n",
+    "        print(f\"  │      ├── {tool}\")\n",
+    "    print(f\"  └── Number of tools: {sample['Annotator Metadata']['Number of tools']}\")\n",
+    "print(\"=\" * 50)"
+   ],
+   "id": "fcf876fdfc43c154",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "### build a vector database based on the metadata.jsonl\n",
+    "# https://python.langchain.com/docs/integrations/vectorstores/supabase/\n",
+    "import os\n",
+    "from dotenv import load_dotenv\n",
+    "from langchain_huggingface import HuggingFaceEmbeddings\n",
+    "from langchain_community.vectorstores import SupabaseVectorStore\n",
+    "from supabase.client import Client, create_client\n",
+    "\n",
+    "\n",
+    "load_dotenv()\n",
+    "embeddings = HuggingFaceEmbeddings(model_name=\"sentence-transformers/all-mpnet-base-v2\") #  dim=768\n",
+    "\n",
+    "supabase_url = os.environ.get(\"SUPABASE_URL\")\n",
+    "supabase_key = os.environ.get(\"SUPABASE_SERVICE_KEY\")\n",
+    "supabase: Client = create_client(supabase_url, supabase_key)"
+   ],
+   "id": "9be49655bf7b1506",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "# wrap the metadata.jsonl's questions and answers into a list of document\n",
+    "from langchain.schema import Document\n",
+    "docs = []\n",
+    "for sample in json_QA:\n",
+    "    content = f\"Question : {sample['Question']}\\n\\nFinal answer : {sample['Final answer']}\"\n",
+    "    doc = {\n",
+    "        \"content\" : content,\n",
+    "        \"metadata\" : { # meatadata的格式必须时source键，否则会报错\n",
+    "            \"source\" : sample['task_id']\n",
+    "        },\n",
+    "        \"embedding\" : embeddings.embed_query(content),\n",
+    "    }\n",
+    "    docs.append(doc)\n",
+    "\n",
+    "# upload the documents to the vector database\n",
+    "try:\n",
+    "    response = (\n",
+    "        supabase.table(\"documents\")\n",
+    "        .insert(docs)\n",
+    "        .execute()\n",
+    "    )\n",
+    "except Exception as exception:\n",
+    "    print(\"Error inserting data into Supabase:\", exception)\n",
+    "\n",
+    "# ALTERNATIVE : Save the documents (a list of dict) into a csv file, and manually upload it to Supabase\n",
+    "# import pandas as pd\n",
+    "# df = pd.DataFrame(docs)\n",
+    "# df.to_csv('supabase_docs.csv', index=False)"
+   ],
+   "id": "3a8f1a5fef11a6a0",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "# add items to vector database\n",
+    "vector_store = SupabaseVectorStore(\n",
+    "    client=supabase,\n",
+    "    embedding= embeddings,\n",
+    "    table_name=\"documents\",\n",
+    "    query_name=\"match_documents_langchain\",\n",
+    ")\n",
+    "retriever = vector_store.as_retriever()"
+   ],
+   "id": "b1b720c17d34d0dc",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "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",
+    "# matched_docs = vector_store.similarity_search(query, 2)\n",
+    "docs = retriever.invoke(query)\n",
+    "docs[0]"
+   ],
+   "id": "49b8420a90e4dd76",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "# list of the tools used in all the samples\n",
+    "from collections import Counter, OrderedDict\n",
+    "\n",
+    "tools = []\n",
+    "for sample in json_QA:\n",
+    "    for tool in sample['Annotator Metadata']['Tools'].split('\\n'):\n",
+    "        tool = tool[2:].strip().lower()\n",
+    "        if tool.startswith(\"(\"):\n",
+    "            tool = tool[11:].strip()\n",
+    "        tools.append(tool)\n",
+    "tools_counter = OrderedDict(Counter(tools))\n",
+    "print(\"List of tools used in all samples:\")\n",
+    "print(\"Total number of tools used:\", len(tools_counter))\n",
+    "for tool, count in tools_counter.items():\n",
+    "    print(f\"  ├── {tool}: {count}\")"
+   ],
+   "id": "6280689c6bf3255e",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "system_prompt = \"\"\"\n",
+    "You are a helpful assistant tasked with answering questions using a set of tools.\n",
+    "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",
+    "You need to provide a step-by-step explanation of how you arrived at the answer.\n",
+    "==========================\n",
+    "Here is a few examples showing you how to answer the question step by step.\n",
+    "\"\"\"\n",
+    "for i, samples in enumerate(random_samples):\n",
+    "    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",
+    "system_prompt += \"\\n==========================\\n\"\n",
+    "system_prompt += \"Now, please answer the following question step by step.\\n\"\n",
+    "\n",
+    "# save the system_prompt to a file\n",
+    "with open('system_prompt.txt', 'w') as f:\n",
+    "    f.write(system_prompt)"
+   ],
+   "id": "d850bbc4a908a98b",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "# load the system prompt from the file\n",
+    "with open('system_prompt.txt', 'r') as f:\n",
+    "    system_prompt = f.read()\n",
+    "print(system_prompt)"
+   ],
+   "id": "85fce16b57c00eaa",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "import os\n",
+    "from dotenv import load_dotenv\n",
+    "from langgraph.graph import MessagesState, START, StateGraph\n",
+    "from langgraph.prebuilt import tools_condition\n",
+    "from langgraph.prebuilt import ToolNode\n",
+    "from langchain_google_genai import ChatGoogleGenerativeAI\n",
+    "from langchain_groq import ChatGroq\n",
+    "from langchain_huggingface import HuggingFaceEmbeddings\n",
+    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
+    "from langchain_community.document_loaders import WikipediaLoader\n",
+    "from langchain_community.document_loaders import ArxivLoader\n",
+    "from langchain_community.vectorstores import SupabaseVectorStore\n",
+    "from langchain.tools.retriever import create_retriever_tool\n",
+    "from langchain_core.messages import HumanMessage, SystemMessage\n",
+    "from langchain_core.tools import tool\n",
+    "from supabase.client import Client, create_client\n",
+    "\n",
+    "# Define the retriever from supabase\n",
+    "load_dotenv()\n",
+    "embeddings = HuggingFaceEmbeddings(model_name=\"sentence-transformers/all-mpnet-base-v2\") #  dim=768\n",
+    "\n",
+    "supabase_url = os.environ.get(\"SUPABASE_URL\")\n",
+    "supabase_key = os.environ.get(\"SUPABASE_SERVICE_KEY\")\n",
+    "supabase: Client = create_client(supabase_url, supabase_key)\n",
+    "vector_store = SupabaseVectorStore(\n",
+    "    client=supabase,\n",
+    "    embedding= embeddings,\n",
+    "    table_name=\"documents\",\n",
+    "    query_name=\"match_documents_langchain\",\n",
+    ")\n",
+    "\n",
+    "question_retrieve_tool = create_retriever_tool(\n",
+    "    vector_store.as_retriever(),\n",
+    "    \"Question Retriever\",\n",
+    "    \"Find similar questions in the vector database for the given question.\",\n",
+    ")\n",
+    "\n",
+    "@tool\n",
+    "def multiply(a: int, b: int) -> int:\n",
+    "    \"\"\"Multiply two numbers.\n",
+    "\n",
+    "    Args:\n",
+    "        a: first int\n",
+    "        b: second int\n",
+    "    \"\"\"\n",
+    "    return a * b\n",
+    "\n",
+    "@tool\n",
+    "def add(a: int, b: int) -> int:\n",
+    "    \"\"\"Add two numbers.\n",
+    "\n",
+    "    Args:\n",
+    "        a: first int\n",
+    "        b: second int\n",
+    "    \"\"\"\n",
+    "    return a + b\n",
+    "\n",
+    "@tool\n",
+    "def subtract(a: int, b: int) -> int:\n",
+    "    \"\"\"Subtract two numbers.\n",
+    "\n",
+    "    Args:\n",
+    "        a: first int\n",
+    "        b: second int\n",
+    "    \"\"\"\n",
+    "    return a - b\n",
+    "\n",
+    "@tool\n",
+    "def divide(a: int, b: int) -> int:\n",
+    "    \"\"\"Divide two numbers.\n",
+    "\n",
+    "    Args:\n",
+    "        a: first int\n",
+    "        b: second int\n",
+    "    \"\"\"\n",
+    "    if b == 0:\n",
+    "        raise ValueError(\"Cannot divide by zero.\")\n",
+    "    return a / b\n",
+    "\n",
+    "@tool\n",
+    "def modulus(a: int, b: int) -> int:\n",
+    "    \"\"\"Get the modulus of two numbers.\n",
+    "\n",
+    "    Args:\n",
+    "        a: first int\n",
+    "        b: second int\n",
+    "    \"\"\"\n",
+    "    return a % b\n",
+    "\n",
+    "@tool\n",
+    "def wiki_search(query: str) -> str:\n",
+    "    \"\"\"Search Wikipedia for a query and return maximum 2 results.\n",
+    "\n",
+    "    Args:\n",
+    "        query: The search query.\"\"\"\n",
+    "    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()\n",
+    "    formatted_search_docs = \"\\n\\n---\\n\\n\".join(\n",
+    "        [\n",
+    "            f'<Document source=\"{doc.metadata[\"source\"]}\" page=\"{doc.metadata.get(\"page\", \"\")}\"/>\\n{doc.page_content}\\n</Document>'\n",
+    "            for doc in search_docs\n",
+    "        ])\n",
+    "    return {\"wiki_results\": formatted_search_docs}\n",
+    "\n",
+    "@tool\n",
+    "def web_search(query: str) -> str:\n",
+    "    \"\"\"Search Tavily for a query and return maximum 3 results.\n",
+    "\n",
+    "    Args:\n",
+    "        query: The search query.\"\"\"\n",
+    "    search_docs = TavilySearchResults(max_results=3).invoke(query=query)\n",
+    "    formatted_search_docs = \"\\n\\n---\\n\\n\".join(\n",
+    "        [\n",
+    "            f'<Document source=\"{doc.metadata[\"source\"]}\" page=\"{doc.metadata.get(\"page\", \"\")}\"/>\\n{doc.page_content}\\n</Document>'\n",
+    "            for doc in search_docs\n",
+    "        ])\n",
+    "    return {\"web_results\": formatted_search_docs}\n",
+    "\n",
+    "@tool\n",
+    "def arvix_search(query: str) -> str:\n",
+    "    \"\"\"Search Arxiv for a query and return maximum 3 result.\n",
+    "\n",
+    "    Args:\n",
+    "        query: The search query.\"\"\"\n",
+    "    search_docs = ArxivLoader(query=query, load_max_docs=3).load()\n",
+    "    formatted_search_docs = \"\\n\\n---\\n\\n\".join(\n",
+    "        [\n",
+    "            f'<Document source=\"{doc.metadata[\"source\"]}\" page=\"{doc.metadata.get(\"page\", \"\")}\"/>\\n{doc.page_content[:1000]}\\n</Document>'\n",
+    "            for doc in search_docs\n",
+    "        ])\n",
+    "    return {\"arvix_results\": formatted_search_docs}\n",
+    "\n",
+    "@tool\n",
+    "def similar_question_search(question: str) -> str:\n",
+    "    \"\"\"Search the vector database for similar questions and return the first results.\n",
+    "\n",
+    "    Args:\n",
+    "        question: the question human provided.\"\"\"\n",
+    "    matched_docs = vector_store.similarity_search(query, 3)\n",
+    "    formatted_search_docs = \"\\n\\n---\\n\\n\".join(\n",
+    "        [\n",
+    "            f'<Document source=\"{doc.metadata[\"source\"]}\" page=\"{doc.metadata.get(\"page\", \"\")}\"/>\\n{doc.page_content[:1000]}\\n</Document>'\n",
+    "            for doc in matched_docs\n",
+    "        ])\n",
+    "    return {\"similar_questions\": formatted_search_docs}\n",
+    "\n",
+    "tools = [\n",
+    "    multiply,\n",
+    "    add,\n",
+    "    subtract,\n",
+    "    divide,\n",
+    "    modulus,\n",
+    "    wiki_search,\n",
+    "    web_search,\n",
+    "    arvix_search,\n",
+    "    question_retrieve_tool\n",
+    "]\n",
+    "\n",
+    "llm = ChatGoogleGenerativeAI(model=\"gemini-2.0-flash\")\n",
+    "# llm = ChatGroq(model=\"qwen-qwq-32b\", temperature=0)\n",
+    "llm_with_tools = llm.bind_tools(tools)"
+   ],
+   "id": "2d049ca6773dd05e",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "# load the system prompt from the file\n",
+    "with open('system_prompt.txt', 'r') as f:\n",
+    "    system_prompt = f.read()\n",
+    "\n",
+    "\n",
+    "# System message\n",
+    "sys_msg = SystemMessage(content=system_prompt)\n",
+    "\n",
+    "# Node\n",
+    "def assistant(state: MessagesState):\n",
+    "    \"\"\"Assistant node\"\"\"\n",
+    "    return {\"messages\": [llm_with_tools.invoke([sys_msg] + state[\"messages\"])]}\n",
+    "\n",
+    "# Build graph\n",
+    "builder = StateGraph(MessagesState)\n",
+    "builder.add_node(\"assistant\", assistant)\n",
+    "builder.add_node(\"tools\", ToolNode(tools))\n",
+    "builder.add_edge(START, \"assistant\")\n",
+    "builder.add_conditional_edges(\n",
+    "    \"assistant\",\n",
+    "    # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools\n",
+    "    # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END\n",
+    "    tools_condition,\n",
+    ")\n",
+    "builder.add_edge(\"tools\", \"assistant\")\n",
+    "\n",
+    "# Compile graph\n",
+    "graph = builder.compile()\n",
+    "\n",
+    "from IPython.display import Image, display\n",
+    "\n",
+    "display(Image(graph.get_graph(xray=True).draw_mermaid_png()))"
+   ],
+   "id": "f1d17cb6cac71a7a",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {
+    "jupyter": {
+     "is_executing": true
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "question = \"What are the EC numbers of the two most commonly used chemicals for the virus testing method in the paper about SPFMV and SPCSV in the Pearl Of Africa from 2016? Return the semicolon-separated numbers in the order of the alphabetized chemicals.\"\n",
+    "messages = [HumanMessage(content=question)]\n",
+    "messages = graph.invoke({\"messages\": messages})\n",
+    "\n",
+    "for m in messages['messages']:\n",
+    "    m.pretty_print()"
+   ],
+   "id": "a1a5be38e7e98476",
+   "outputs": [],
+   "execution_count": null
+  }
+ ],
+ "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
+}