Update agent.py

agent.py

@@ -1,23 +1,70 @@
+import os
+from dotenv import load_dotenv
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import tools_condition, ToolNode
 from langchain_groq import ChatGroq
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_community.vectorstores import SupabaseVectorStore
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
-from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_community.document_loaders import WikipediaLoader
-from langgraph.graph import START, StateGraph, MessagesState
-from langgraph.prebuilt import tools_condition
-from langgraph.prebuilt import ToolNode
-from dotenv import load_dotenv
-import os
-from supabase.client import Client, create_client
-import json
-from langchain.schema import Document
+from langchain.tools.retriever import create_retriever_tool
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.document_loaders import ArxivLoader
-from langchain_community.vectorstores import SupabaseVectorStore
-
+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:
@@ -62,19 +109,16 @@ def arvix_search(query: str) -> str:
     return {"arvix_results": formatted_search_docs}
 
 
-tools = [
-    wiki_search,
-    web_search,
-    arvix_search,
-]
 
 # 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)
 
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2"
-
+# 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"))
@@ -84,14 +128,43 @@ vector_store = SupabaseVectorStore(
     table_name="documents",
     query_name="match_documents",
 )
+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.",
+)
 
-# System message
-sys_msg = SystemMessage(content=system_prompt)
 
 
+tools = [
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulus,
+    wiki_search,
+    web_search,
+    arvix_search,
+]
+
 # Build graph function
-def build_graph():
-    llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
+def build_graph(provider: str = "groq"):
+    """Build the graph"""
+    # Load environment variables from .env file
+    if 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, 'groq' or 'huggingface'.")
+    # Bind tools to LLM
     llm_with_tools = llm.bind_tools(tools)
 
     # Node
@@ -106,7 +179,7 @@ def build_graph():
             content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
         )
         return {"messages": [sys_msg] + state["messages"] + [example_msg]}
-    
+
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
@@ -121,73 +194,14 @@ def build_graph():
 
     # Compile graph
     return builder.compile()
-    
-
-def load_documents():
-
-    embeddings = AzureOpenAIEmbeddings(model="text-embedding-3-small")
-    with open('metadata.jsonl', 'r') as jsonl_file:
-        json_list = list(jsonl_file)
-
-    json_QA = []
-    for json_str in json_list:
-        json_data = json.loads(json_str)
-        json_QA.append(json_data)
-
-    supabase_url = os.environ.get("SUPABASE_URL")
-    supabase_key = os.environ.get("SUPABASE_SERVICE_KEY")
-    supabase: Client = create_client(supabase_url, supabase_key)
-
-    docs: list[Document] = []
-    for sample in json_QA:
-        content = f"Question : {sample['Question']}\n\nFinal answer : {sample['Final answer']}"
-        doc = {
-            "content" : content,
-            "metadata" : { 
-                "source" : sample['task_id']
-            },
-            "embedding" : embeddings.embed_query(content),
-        }
-        docs.append(doc)
-
-    # upload the documents to the vector database
-    try:
-        response = (
-            supabase.table("documents")
-            .insert(docs)
-            .execute()
-        )
-    except Exception as exception:
-        print("Error inserting data into Supabase:", exception)
-
-
-def search_documents() -> list[Document]:
-    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?"
-    vector_store = SupabaseVectorStore(
-        client=supabase,
-        embedding= embeddings,
-        table_name="documents",
-        query_name="match_documents",
-    )
-    retriever = vector_store.as_retriever()
-    docs = retriever.get_relevant_documents(query)
-    return docs[0]
-
-
-
-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()
-    # Run the graph
-    messages = [HumanMessage(content=question)]
-    messages = graph.invoke({"messages": messages})
-    for m in messages["messages"]:
-        m.pretty_print()
-
-    # load_documents()
-    # search_documents()
-
-    
-
 
+# 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")
+#     # Run the graph
+#     messages = [HumanMessage(content=question)]
+#     messages = graph.invoke({"messages": messages})
+#     for m in messages["messages"]:
+#         m.pretty_print()