Update agent.py

agent.py

@@ -1,4 +1,4 @@
-"""LangGraph Agent (No Supabase)"""
+"""LangGraph Agent (GROQ version without Supabase)"""
 import os
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import tools_condition, ToolNode
@@ -12,57 +12,63 @@ from langchain_core.tools import tool
 
 @tool
 def multiply(a: int, b: int) -> int:
-    """Multiply two integers and return the result."""
+    """Multiply two numbers."""
     return a * b
 
 @tool
 def add(a: int, b: int) -> int:
-    """Add two integers and return the result."""
+    """Add two numbers."""
     return a + b
 
 @tool
 def subtract(a: int, b: int) -> int:
-    """Subtract b from a and return the result."""
+    """Subtract second number from the first."""
     return a - b
 
 @tool
 def divide(a: int, b: int) -> float:
-    """Divide a by b and return the result. Raises an error if b is zero."""
+    """Divide two numbers."""
     if b == 0:
         raise ValueError("Cannot divide by zero.")
     return a / b
 
 @tool
 def modulus(a: int, b: int) -> int:
-    """Return the modulus (remainder) of a divided by b."""
+    """Get the modulus (remainder) of two numbers."""
     return a % b
 
 @tool
 def wiki_search(query: str) -> dict:
     """Search Wikipedia for a query and return up to 2 results."""
     search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    results = "\n\n---\n\n".join(
-        f"<Document>\n{doc.page_content}\n</Document>" for doc in search_docs
-    )
-    return {"wiki_results": results}
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"wiki_results": formatted_search_docs}
 
 @tool
 def web_search(query: str) -> dict:
-    """Search the web via Tavily and return up to 3 results."""
+    """Search Tavily for a query and return up to 3 results."""
     search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-    results = "\n\n---\n\n".join(
-        f"<Document>\n{doc.page_content}\n</Document>" for doc in search_docs
-    )
-    return {"web_results": results}
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in search_docs
+        ])
+    return {"web_results": formatted_search_docs}
 
 @tool
 def arvix_search(query: str) -> dict:
     """Search Arxiv and return up to 3 truncated results."""
     search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-    results = "\n\n---\n\n".join(
-        f"<Document>\n{doc.page_content[:500]}\n</Document>" for doc in search_docs
-    )
-    return {"arvix_results": results}
+    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 search_docs
+        ])
+    return {"arvix_results": formatted_search_docs}
 
 # Load system prompt
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
@@ -76,14 +82,17 @@ tools = [
 ]
 
 def build_graph(provider: str = "groq"):
-    """Build the LangGraph agent with selected LLM provider."""
+    """Build the LangGraph agent using specified LLM provider."""
+
     if provider == "google":
         llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+
     elif provider == "groq":
-        groq_api_key = os.environ.get("GROQ_API_KEY")
+        groq_api_key = os.getenv("GROQ_API_KEY")
         if not groq_api_key:
-            raise ValueError("GROQ_API_KEY is not set in the environment.")
+            raise ValueError("GROQ_API_KEY environment variable not set.")
         llm = ChatGroq(model="qwen-qwq-32b", temperature=0, api_key=groq_api_key)
+
     elif provider == "huggingface":
         llm = ChatHuggingFace(
             llm=HuggingFaceEndpoint(
@@ -91,8 +100,9 @@ def build_graph(provider: str = "groq"):
                 temperature=0,
             )
         )
+
     else:
-        raise ValueError("Invalid provider: choose 'google', 'groq' or 'huggingface'.")
+        raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
 
     llm_with_tools = llm.bind_tools(tools)
 
@@ -108,10 +118,10 @@ def build_graph(provider: str = "groq"):
 
     return builder.compile()
 
+# For testing purposes
 if __name__ == "__main__":
-    from langchain_core.messages import HumanMessage
-    question = "What is the capital of France and its population?"
-    graph = build_graph()
+    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
+    graph = build_graph(provider="groq")
     messages = [HumanMessage(content=question)]
     result = graph.invoke({"messages": messages})
     for msg in result["messages"]: