Update agent.py

agent.py

@@ -2,51 +2,79 @@
 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 langgraph.prebuilt import tools_condition, 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.document_loaders import WikipediaLoader, ArxivLoader
 from langchain_community.vectorstores import SupabaseVectorStore
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain.tools.retriever import create_retriever_tool
-from supabase.client import Client, create_client
+from supabase.client import create_client
+import re
 
 load_dotenv()
 
 # ------------------ Arithmetic Tools ------------------
+
 @tool
 def multiply(a: int, b: int) -> str:
-    """Multiply two integers and return the product as a string."""
+    """
+    Multiply two integers and return the result as a string.
+
+    Args:
+        a (int): The first integer.
+        b (int): The second integer.
+
+    Returns:
+        str: The product of a and b, as a string.
+    """
     return str(a * b)
 
 
 @tool
 def add(a: int, b: int) -> str:
-    """Add two integers and return the sum as a string."""
+    """
+    Add two integers and return the result as a string.
+
+    Args:
+        a (int): The first integer.
+        b (int): The second integer.
+
+    Returns:
+        str: The sum of a and b, as a string.
+    """
     return str(a + b)
 
+
 @tool
 def subtract(a: int, b: int) -> str:
-    """Subtract two integers and return the difference as a string."""
+    """
+    Subtract one integer from another and return the result as a string.
+
+    Args:
+        a (int): The minuend.
+        b (int): The subtrahend.
+
+    Returns:
+        str: The difference (a - b), as a string.
+    """
     return str(a - b)
 
+
 @tool
 def divide(a: int, b: int) -> str:
     """
-    Divide two integers and return the result as a string.
+    Divide one integer by another and return the result as a string.
 
     Args:
-        a: The numerator (dividend).
-        b: The denominator (divisor). Must not be zero.
+        a (int): The numerator.
+        b (int): The denominator. Must not be zero.
 
     Returns:
-        A string representation of the division result, or an error message if b is zero.
+        str: The result of the division (a / b), as a string. Returns an error message if b is zero.
     """
     if b == 0:
         return "Error: Cannot divide by zero."
@@ -59,26 +87,27 @@ def modulus(a: int, b: int) -> str:
     Compute the modulus (remainder) of two integers and return the result as a string.
 
     Args:
-        a: The numerator.
-        b: The denominator.
+        a (int): The numerator.
+        b (int): The denominator.
 
     Returns:
-        A string representation of the remainder when a is divided by b.
+        str: The remainder when a is divided by b, as a string.
     """
     return str(a % b)
 
 
 # ------------------ Retrieval Tools ------------------
+
 @tool
 def wiki_search(query: str) -> str:
     """
-    Search Wikipedia for a given query and return the content of up to two matching articles.
+    Search Wikipedia for a given query and return text from up to two matching articles.
 
     Args:
-        query: A string query to search on Wikipedia.
+        query (str): A string query to search on Wikipedia.
 
     Returns:
-        A string containing the content from up to two relevant Wikipedia articles, separated by dividers.
+        str: Combined content from up to two relevant articles, separated by dividers.
     """
     docs = WikipediaLoader(query=query, load_max_docs=2).load()
     return "\n\n---\n\n".join(doc.page_content for doc in docs)
@@ -90,10 +119,10 @@ def web_search(query: str) -> str:
     Perform a web search using Tavily and return content from the top three results.
 
     Args:
-        query: A string query representing the web search topic.
+        query (str): A string representing the web search topic.
 
     Returns:
-        A string of up to three relevant result contents, separated by dividers.
+        str: Combined content from up to three top results, separated by dividers.
     """
     docs = TavilySearchResults(max_results=3).invoke(query)
     return "\n\n---\n\n".join(doc.page_content for doc in docs)
@@ -102,37 +131,35 @@ def web_search(query: str) -> str:
 @tool
 def arvix_search(query: str) -> str:
     """
-    Search arXiv for academic papers matching the query and return excerpts from up to three results.
+    Search arXiv for academic papers related to the query and return excerpts.
 
     Args:
-        query: A string query to search on arXiv.
+        query (str): The search query string.
 
     Returns:
-        A string containing excerpts (first 1000 characters) from up to three relevant papers.
+        str: Excerpts (up to 1000 characters each) from up to three relevant arXiv papers, separated by dividers.
     """
     docs = ArxivLoader(query=query, load_max_docs=3).load()
     return "\n\n---\n\n".join(doc.page_content[:1000] for doc in docs)
 
 
-# ------------------ System Setup ------------------
+
+# ------------------ System Prompt ------------------
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
     system_prompt = f.read().strip()
-sys_msg = SystemMessage(content=system_prompt)
-
 
-# Load environment variables
-url = os.environ["SUPABASE_URL"]
-key = os.environ["SUPABASE_SERVICE_KEY"]
+# ------------------ Supabase Setup ------------------
+url = os.environ["SUPABASE_URL"].strip()
+key = os.environ["SUPABASE_SERVICE_KEY"].strip()
 client = create_client(url, key)
 
-# Create embedding model
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
-# Sample documents to insert
-docs = [
-    {"content": "Newton's First Law states that an object in motion stays in motion unless acted upon."},
-    {"content": "LangChain enables developers to build context-aware agents using LLMs and tools."},
-    {"content": "Supabase is an open-source alternative to Firebase built on PostgreSQL."}
+# Embed improved QA docs
+qa_examples = [
+    {"content": "Q: What is the capital of Vietnam?\nA: FINAL ANSWER: Hanoi"},
+    {"content": "Q: Alphabetize: lettuce, broccoli, basil\nA: FINAL ANSWER: basil,broccoli,lettuce"},
+    {"content": "Q: What is 42 multiplied by 8?\nA: FINAL ANSWER: three hundred thirty six"},
 ]
 vector_store = SupabaseVectorStore(
     client=client,
@@ -140,10 +167,8 @@ vector_store = SupabaseVectorStore(
     table_name="documents",
     query_name="match_documents_langchain"
 )
-texts = [doc["content"] for doc in docs]
-vector_store.add_texts(texts)
-
-print("✅ Documents successfully embedded and stored.")
+vector_store.add_texts([doc["content"] for doc in qa_examples])
+print("✅ QA documents embedded into Supabase.")
 
 retriever_tool = create_retriever_tool(
     retriever=vector_store.as_retriever(),
@@ -154,16 +179,21 @@ retriever_tool = create_retriever_tool(
 tools = [multiply, add, subtract, divide, modulus, wiki_search, web_search, arvix_search]
 
 # ------------------ Build Agent Graph ------------------
+class VerboseToolNode(ToolNode):
+    def invoke(self, state):
+        print("🔧 ToolNode evaluating:", [m.content for m in state["messages"]])
+        return super().invoke(state)
+
 def build_graph(provider: str = "groq"):
     if provider == "google":
-        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0.3)
     elif provider == "groq":
-        llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
+        llm = ChatGroq(model="qwen-qwq-32b", temperature=0.3)
     elif provider == "huggingface":
         llm = ChatHuggingFace(
             llm=HuggingFaceEndpoint(
                 url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
-                temperature=0
+                temperature=0.3
             )
         )
     else:
@@ -172,48 +202,38 @@ def build_graph(provider: str = "groq"):
     llm_with_tools = llm.bind_tools(tools)
 
     def retriever(state: MessagesState):
-        similar = vector_store.similarity_search(state["messages"][0].content)
+        query = state["messages"][0].content
+        similar = vector_store.similarity_search_with_score(query)
+        threshold = 0.7
         examples = [
             HumanMessage(content=f"Similar QA:\n{doc.page_content}")
-            for doc in similar[:2]
+            for doc, score in similar if score >= threshold
         ]
-        return {"messages": [sys_msg] + state["messages"] + examples}
-
+        return {"messages": state["messages"] + examples}
 
     def assistant(state: MessagesState):
         try:
-            # Always prepend the system message
-            system_msg = SystemMessage(content=system_prompt.strip())
-            messages = [system_msg] + state["messages"]
-
+            messages = [SystemMessage(content=system_prompt.strip())] + state["messages"]
             result = llm_with_tools.invoke(messages)
-            print("🤖 Raw LLM result:", repr(result))
-
             raw_output = result.content.strip()
+            print("🤖 Raw LLM output:", repr(raw_output))
 
-            # Extract FINAL ANSWER using regex
-            import re
             match = re.search(r"FINAL ANSWER:\s*(.+)", raw_output, re.IGNORECASE)
             if match:
-                final_answer = match.group(1).strip()
-                final_output = f"FINAL ANSWER: {final_answer}"
+                final_output = f"FINAL ANSWER: {match.group(1).strip()}"
             else:
-                print("⚠️ 'FINAL ANSWER:' prefix not found. Using fallback.")
-                final_output = "FINAL ANSWER: Unable to determine answer"
+                print("⚠️ 'FINAL ANSWER:' not found. Raw content will be used as fallback.")
+                final_output = f"FINAL ANSWER: {raw_output or 'Unable to determine answer'}"
 
             return {"messages": [HumanMessage(content=final_output)]}
-
         except Exception as e:
             print(f"🔥 Error in assistant node: {e}")
             return {"messages": [HumanMessage(content=f"FINAL ANSWER: AGENT ERROR: {e}")]}
 
-
-
-
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
+    builder.add_node("tools", VerboseToolNode(tools))
     builder.add_edge(START, "retriever")
     builder.add_edge("retriever", "assistant")
     builder.add_conditional_edges("assistant", tools_condition)
@@ -228,3 +248,4 @@ if __name__ == "__main__":
     messages = [HumanMessage(content=question)]
     result = graph.invoke({"messages": messages})
     print(result["messages"][-1].content)
+