add retriever tool

agent.py

@@ -1,15 +1,22 @@
 """LangGraph Agent"""
-import dotenv
-from langgraph.graph import MessagesState
-from langgraph.graph import START, StateGraph
+import os
+from dotenv import load_dotenv
+from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import tools_condition
 from langgraph.prebuilt import ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_groq import ChatGroq
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
 from langchain_community.document_loaders import ArxivLoader
+from langchain_community.vectorstores import SupabaseVectorStore
+from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
-from langchain_core.messages import SystemMessage
+from langchain.tools.retriever import create_retriever_tool
+from supabase.client import Client, create_client
+
+load_dotenv()
 
 @tool
 def multiply(a: int, b: int) -> int:
@@ -105,6 +112,25 @@ def arvix_search(query: str) -> str:
         ])
     return {"arvix_results": formatted_search_docs}
 
+# build a retriever tool
+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"))
+vector_store = SupabaseVectorStore(
+    client=supabase,
+    embedding= embeddings,
+    table_name="documents",
+    query_name="match_documents_langchain",
+)
+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,
     add,
@@ -114,14 +140,9 @@ tools = [
     wiki_search,
     web_search,
     arvix_search,
+    question_retrieve_tool
 ]
 
-
-# Load environment variables from .env file
-dotenv.load_dotenv()
-llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash")
-llm_with_tools = llm.bind_tools(tools)
-
 # load the system prompt from the file
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
     system_prompt = f.read()
@@ -129,14 +150,35 @@ with open("system_prompt.txt", "r", encoding="utf-8") as f:
 # System message
 sys_msg = SystemMessage(content=system_prompt)
 
-# Node
-def assistant(state: MessagesState):
-    """Assistant node"""
-    return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
+
 
 # Build graph function
-def build_graph():
+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="gemma2-9b-it", temperature=0) # optional : qwen-qwq-32b
+    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
+    llm_with_tools = llm.bind_tools(tools)
+
+    # Node
+    def assistant(state: MessagesState):
+        """Assistant node"""
+        return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
 
     builder = StateGraph(MessagesState)
     builder.add_node("assistant", assistant)
@@ -150,3 +192,16 @@ def build_graph():
 
     # Compile graph
     return builder.compile()
+
+# test
+if __name__ == "__main__":
+    question = "Who nominated the only Featured Article on English Wikipedia about a dinosaur that was promoted in November 2016?"
+
+    # Build the graph
+    graph = build_graph(provider="groq")
+    # Run the graph
+    messages = [HumanMessage(content=question)]
+    messages = graph.invoke({"messages": messages})
+    answer = messages[-1].content
+    print(f"Question: {question}")
+    print(f"{answer}")

requirements.txt

@@ -4,9 +4,15 @@ langchain
 langchain-community
 langchain-core
 langchain-google-genai
+langchain-huggingface
+langchain-groq
 langchain-tavily
+langchain-chroma
 langgraph
+huggingface_hub
+supabase
 arxiv
 pymupdf
 wikipedia
-dotenv
+pgvector
+python-dotenv

system_prompt.txt

@@ -33,4 +33,6 @@ Tools:
 Final Answer: Rd5
 ==========================
 
-Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+Now, I will ask you a question. Report your thoughts, and finish your answer with the following template: 
+FINAL ANSWER: [YOUR FINAL ANSWER]. 
+YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.

test.ipynb

@@ -1,8 +1,16 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "d0cc4adf",
+   "metadata": {},
+   "source": [
+    "### Question data"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 50,
+   "execution_count": 2,
    "id": "14e3f417",
    "metadata": {},
    "outputs": [],
@@ -21,7 +29,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 51,
+   "execution_count": 3,
    "id": "5e2da6fc",
    "metadata": {},
    "outputs": [
@@ -30,45 +38,28 @@
      "output_type": "stream",
      "text": [
       "==================================================\n",
-      "Task ID: f0f46385-fc03-4599-b5d3-f56496c3e69f\n",
-      "Question: In terms of geographical distance between capital cities, which 2 countries are the furthest from each other within the ASEAN bloc according to wikipedia? Answer using a comma separated list, ordering the countries by alphabetical order.\n",
+      "Task ID: ed58682d-bc52-4baa-9eb0-4eb81e1edacc\n",
+      "Question: What is the last word before the second chorus of the King of Pop's fifth single from his sixth studio album?\n",
       "Level: 2\n",
-      "Final Answer: Indonesia, Myanmar\n",
+      "Final Answer: stare\n",
       "Annotator Metadata: \n",
       "  ├── Steps: \n",
-      "  │      ├── 1. Search the web for \"ASEAN bloc\".\n",
-      "  │      ├── 2. Click the Wikipedia result for the ASEAN Free Trade Area.\n",
-      "  │      ├── 3. Scroll down to find the list of member states.\n",
-      "  │      ├── 4. Click into the Wikipedia pages for each member state, and note its capital.\n",
-      "  │      ├── 5. Search the web for the distance between the first two capitals. The results give travel distance, not geographic distance, which might affect the answer.\n",
-      "  │      ├── 6. Thinking it might be faster to judge the distance by looking at a map, search the web for \"ASEAN bloc\" and click into the images tab.\n",
-      "  │      ├── 7. View a map of the member countries. Since they're clustered together in an arrangement that's not very linear, it's difficult to judge distances by eye.\n",
-      "  │      ├── 8. Return to the Wikipedia page for each country. Click the GPS coordinates for each capital to get the coordinates in decimal notation.\n",
-      "  │      ├── 9. Place all these coordinates into a spreadsheet.\n",
-      "  │      ├── 10. Write formulas to calculate the distance between each capital.\n",
-      "  │      ├── 11. Write formula to get the largest distance value in the spreadsheet.\n",
-      "  │      ├── 12. Note which two capitals that value corresponds to: Jakarta and Naypyidaw.\n",
-      "  │      ├── 13. Return to the Wikipedia pages to see which countries those respective capitals belong to: Indonesia, Myanmar.\n",
-      "  ├── Number of steps: 13\n",
-      "  ├── How long did this take?: 45 minutes\n",
+      "  │      ├── 1. Google searched \"King of Pop\".\n",
+      "  │      ├── 2. Clicked on Michael Jackson's Wikipedia.\n",
+      "  │      ├── 3. Scrolled down to \"Discography\".\n",
+      "  │      ├── 4. Clicked on the sixth album, \"Thriller\".\n",
+      "  │      ├── 5. Looked under \"Singles from Thriller\".\n",
+      "  │      ├── 6. Clicked on the fifth single, \"Human Nature\".\n",
+      "  │      ├── 7. Google searched \"Human Nature Michael Jackson Lyrics\".\n",
+      "  │      ├── 8. Looked at the opening result with full lyrics sourced by Musixmatch.\n",
+      "  │      ├── 9. Looked for repeating lyrics to determine the chorus.\n",
+      "  │      ├── 10. Determined the chorus begins with \"If they say\" and ends with \"Does he do me that way?\"\n",
+      "  │      ├── 11. Found the second instance of the chorus within the lyrics.\n",
+      "  │      ├── 12. Noted the last word before the second chorus - \"stare\".\n",
+      "  ├── Number of steps: 12\n",
+      "  ├── How long did this take?: 20 minutes\n",
       "  ├── Tools:\n",
-      "  │      ├── 1. Search engine\n",
-      "  │      ├── 2. Web browser\n",
-      "  │      ├── 3. Microsoft Excel / Google Sheets\n",
-      "  └── Number of tools: 3\n",
-      "==================================================\n",
-      "Task ID: cca530fc-4052-43b2-b130-b30968d8aa44\n",
-      "Question: Review the chess position provided in the image. It is black's turn. Provide the correct next move for black which guarantees a win. Please provide your response in algebraic notation.\n",
-      "Level: 1\n",
-      "Final Answer: Rd5\n",
-      "Annotator Metadata: \n",
-      "  ├── Steps: \n",
-      "  │      ├── Step 1: Evaluate the position of the pieces in the chess position\n",
-      "  │      ├── Step 2: Report the best move available for black: \"Rd5\"\n",
-      "  ├── Number of steps: 2\n",
-      "  ├── How long did this take?: 10 minutes\n",
-      "  ├── Tools:\n",
-      "  │      ├── 1. Image recognition tools\n",
+      "  │      ├── Web Browser\n",
       "  └── Number of tools: 1\n",
       "==================================================\n"
      ]
@@ -80,7 +71,7 @@
     "\n",
     "import random\n",
     "# random.seed(42)\n",
-    "random_samples = random.sample(json_QA, 2)\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",
@@ -100,6 +91,118 @@
     "print(\"=\" * 50)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 48,
+   "id": "4bb02420",
+   "metadata": {},
+   "outputs": [],
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a070b955",
+   "metadata": {},
+   "outputs": [],
+   "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)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 54,
+   "id": "77fb9dbb",
+   "metadata": {},
+   "outputs": [],
+   "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()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 55,
+   "id": "12a05971",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
+      "To disable this warning, you can either:\n",
+      "\t- Avoid using `tokenizers` before the fork if possible\n",
+      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
+     ]
+    },
+    {
+     "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": 55,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "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]"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 31,
@@ -216,6 +319,14 @@
     "    print(f\"  ├── {tool}: {count}\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "5efee12a",
+   "metadata": {},
+   "source": [
+    "#### Graph"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 55,
@@ -301,22 +412,45 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 46,
+   "execution_count": null,
    "id": "42fde0f8",
    "metadata": {},
    "outputs": [],
    "source": [
     "import dotenv\n",
-    "from langgraph.graph import MessagesState\n",
-    "from langchain_core.messages import HumanMessage, SystemMessage\n",
-    "from langgraph.graph import START, StateGraph\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_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",
@@ -412,6 +546,20 @@
     "        ])\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",
@@ -421,11 +569,9 @@
     "    wiki_search,\n",
     "    web_search,\n",
     "    arvix_search,\n",
+    "    question_retrieve_tool\n",
     "]\n",
     "\n",
-    "\n",
-    "# Load environment variables from .env file\n",
-    "dotenv.load_dotenv()\n",
     "llm = ChatGoogleGenerativeAI(model=\"gemini-2.0-flash\")\n",
     "llm_with_tools = llm.bind_tools(tools)"
    ]
@@ -489,6 +635,29 @@
     "\n",
     "display(Image(graph.get_graph(xray=True).draw_mermaid_png()))"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5987d58c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "question = \"\"\n",
+    "messages = [HumanMessage(content=question)]\n",
+    "messages = graph.invoke({\"messages\": messages})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "330cbf17",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for m in messages['messages']:\n",
+    "    m.pretty_print()"
+   ]
   }
  ],
  "metadata": {