Update agent.py

agent.py

@@ -20,68 +20,76 @@ load_dotenv()
 
 @tool
 def multiply(a: int, b: int) -> int:
+    """Multiply two integers and return the product."""
     return a * b
 
 @tool
 def add(a: int, b: int) -> int:
+    """Add two integers and return the sum."""
     return a + b
 
 @tool
 def subtract(a: int, b: int) -> int:
+    """Subtract second integer from first and return the difference."""
     return a - b
 
 @tool
 def divide(a: int, b: int) -> float:
+    """Divide first integer by second and return the quotient. Raises error if divisor is zero."""
     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 first integer divided by second."""
     return a % b
 
 @tool
-def wiki_search(query: str) -> str:
+def wiki_search(query: str) -> dict:
+    """Search Wikipedia for a query and return formatted top 2 results."""
     search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    formatted = "\n\n---\n\n".join(
+    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}
+    return {"wiki_results": formatted_search_docs}
 
 @tool
-def web_search(query: str) -> str:
+def web_search(query: str) -> dict:
+    """Search the web via Tavily and return formatted top 3 results."""
     search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-    formatted = "\n\n---\n\n".join(
+    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}
+    return {"web_results": formatted_search_docs}
 
 @tool
-def arvix_search(query: str) -> str:
+def arvix_search(query: str) -> dict:
+    """Search Arxiv for a query and return formatted top 3 results (truncated content)."""
     search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-    formatted = "\n\n---\n\n".join(
+    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}
+    return {"arvix_results": formatted_search_docs}
 
-
-# Load system prompt
+# Load the system prompt from file
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
     system_prompt = f.read()
 
 sys_msg = SystemMessage(content=system_prompt)
 
-# Setup vector store retriever
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+# Build 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"))
+    os.environ.get("SUPABASE_URL"),
+    os.environ.get("SUPABASE_SERVICE_KEY"),
+)
 vector_store = SupabaseVectorStore(
     client=supabase,
     embedding=embeddings,
@@ -106,6 +114,7 @@ tools = [
 ]
 
 def build_graph(provider: str = "groq"):
+    """Build the LangGraph agent graph with the specified provider."""
     if provider == "google":
         llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
     elif provider == "groq":
@@ -119,13 +128,14 @@ def build_graph(provider: str = "groq"):
         )
     else:
         raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
-
     llm_with_tools = llm.bind_tools(tools)
 
     def assistant(state: MessagesState):
+        """Assistant node to process messages with LLM and tools."""
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
     def retriever(state: MessagesState):
+        """Retriever node to find similar questions from vector store."""
         similar_question = vector_store.similarity_search(state["messages"][0].content)
         example_msg = HumanMessage(
             content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
@@ -143,7 +153,6 @@ def build_graph(provider: str = "groq"):
 
     return builder.compile()
 
-
 if __name__ == "__main__":
     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")