updated

agent.py

@@ -1,25 +1,74 @@
+"""LangGraph Agent"""
 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_nvidia_ai_endpoints import ChatNVIDIA
-
+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.tools.retriever import create_retriever_tool
-from langchain_chroma import Chroma
-from langchain_core.messages import SystemMessage, HumanMessage
-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:
+    """Multiply two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> int:
+    """Divide two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers.
+    
+    Args:
+        a: first int
+        b: second int
+    """
+    return a % b
+
 @tool
 def wiki_search(query: str) -> str:
     """Search Wikipedia for a query and return maximum 2 results.
@@ -64,38 +113,43 @@ def arvix_search(query: str) -> str:
 
 
 
-
 # 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
-
-vector_store = Chroma(
-    collection_name="example_collection",
-    embedding_function=embeddings,
-    persist_directory="./chroma_langchain_db_1",  # Where to save data locally, remove if not necessary
+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",
 )
-
-
-# Assign the result to a new variable name, like 'question_retriever_tool'
-question_retriever_tool = create_retriever_tool(
+create_retriever_tool = create_retriever_tool(
     retriever=vector_store.as_retriever(),
-    name="question_search",  # Changed name to be valid
+    name="Question Search",
     description="A tool to retrieve similar questions from a vector store.",
 )
 
+
+
 tools = [
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulus,
     wiki_search,
     web_search,
     arvix_search,
-    question_retriever_tool,
 ]
 
-
-
 # Build graph function
 def build_graph(provider: str = "groq"):
     """Build the graph"""
@@ -103,9 +157,9 @@ def build_graph(provider: str = "groq"):
     if provider == "google":
         # Google Gemini
         llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-    elif provider == "Nvidia":
-        
-        llm = ChatNVIDIA(model="meta/llama-3.1-70b-instruct", 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(
@@ -151,7 +205,7 @@ def build_graph(provider: str = "groq"):
 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="Nvidia")
+    graph = build_graph(provider="groq")
     # Run the graph
     messages = [HumanMessage(content=question)]
     messages = graph.invoke({"messages": messages})