Update agent.py

agent.py

@@ -1,132 +1,236 @@
 import os
 from dotenv import load_dotenv
+from supabase.client import create_client
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import ToolNode, tools_condition
 from langchain_core.tools import tool
 from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
-from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
+from langchain_huggingface import (
+    ChatHuggingFace,
+    HuggingFaceEndpoint,
+    HuggingFaceEmbeddings,
+)
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
 from langchain_community.vectorstores import SupabaseVectorStore
 from langchain.tools.retriever import create_retriever_tool
-from supabase.client import create_client
 
 load_dotenv()
 
-# --- System Prompt Loader ---
-def load_system_prompt(path="system_prompt.txt") -> SystemMessage:
+
+def load_system_prompt(path: str = "system_prompt.txt") -> SystemMessage:
+    """
+    Load system prompt from a file, fallback to a default if missing.
+
+    Args:
+        path: File path to the system prompt.
+
+    Returns:
+        SystemMessage containing the loaded or default prompt.
+    """
     try:
         with open(path, encoding="utf-8") as f:
-            return SystemMessage(content=f.read())
+            content = f.read()
     except FileNotFoundError:
-        return SystemMessage(content="You are a helpful assistant.")
+        content = "You are a helpful assistant."
+    return SystemMessage(content=content)
 
-sys_msg = load_system_prompt()
 
-# --- Math Tools Factory ---
-def math_tool(fn):
-    return tool(fn)
+def math_tool(func):
+    """
+    Wrap a Python function as a LangChain tool.
+
+    Args:
+        func: Callable to wrap.
+
+    Returns:
+        A LangChain tool.
+    """
+    return tool(func)
+
 
 @math_tool
-def add(a: int, b: int) -> int:        return a + b
+def add(a: int, b: int) -> int:
+    """Return a + b."""
+    return a + b
+
+
 @math_tool
-def subtract(a: int, b: int) -> int:   return a - b
+def subtract(a: int, b: int) -> int:
+    """Return a - b."""
+    return a - b
+
+
 @math_tool
-def multiply(a: int, b: int) -> int:   return a * b
+def multiply(a: int, b: int) -> int:
+    """Return a * b."""
+    return a * b
+
+
 @math_tool
 def divide(a: int, b: int) -> float:
-    if b == 0: raise ValueError("Cannot divide by zero.")
+    """
+    Return a / b.
+
+    Raises:
+        ValueError: If b is zero.
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
     return a / b
 
+
 @math_tool
-def modulus(a: int, b: int) -> int:    return a % b
+def modulus(a: int, b: int) -> int:
+    """Return a % b."""
+    return a % b
+
 
-# --- Document Formatting Helper ---
 def format_docs(docs, key: str, max_chars: int = None) -> dict:
-    content = "\n\n---\n\n".join(
-        f'<Document source="{d.metadata.get("source","")}" page="{d.metadata.get("page","")}" />\n'
-        f'{d.page_content[:max_chars] if max_chars else d.page_content}\n</Document>'
-        for d in docs
-    )
-    return {key: content}
+    """
+    Convert document list into labeled XML-style chunks.
+
+    Args:
+        docs: Iterable of Document objects.
+        key: Dict key for formatted results.
+        max_chars: Optionally truncate content.
+
+    Returns:
+        {key: formatted_string}
+    """
+    entries = []
+    for d in docs:
+        content = d.page_content if max_chars is None else d.page_content[:max_chars]
+        entries.append(
+            f'<Document source="{d.metadata.get("source","")}" page="{d.metadata.get("page","")}">\n'
+            f"{content}\n</Document>"
+        )
+    return {key: "\n\n---\n\n".join(entries)}
+
 
-# --- Info Tools ---
 @tool
 def wiki_search(query: str) -> dict:
+    """Search Wikipedia (2 docs) and format results."""
     docs = WikipediaLoader(query=query, load_max_docs=2).load()
     return format_docs(docs, "wiki_results")
 
+
 @tool
 def web_search(query: str) -> dict:
+    """Search the web via Tavily (3 docs) and format results."""
     docs = TavilySearchResults(max_results=3).invoke(query=query)
     return format_docs(docs, "web_results")
 
+
 @tool
-def arvix_search(query: str) -> dict:
+def arxiv_search(query: str) -> dict:
+    """Search ArXiv (3 docs) and format results (truncate to 1k chars)."""
     docs = ArxivLoader(query=query, load_max_docs=3).load()
-    return format_docs(docs, "arvix_results", max_chars=1000)
+    return format_docs(docs, "arxiv_results", max_chars=1000)
+
 
-# --- Vector Retriever Setup ---
 def build_vector_retriever():
-    embed_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-    supa = create_client(os.getenv("SUPABASE_URL"), os.getenv("SUPABASE_SERVICE_KEY"))
-    vs = SupabaseVectorStore(
+    """
+    Create and return a Supabase-based vector retriever.
+
+    Returns:
+        Retriever for semantic similarity queries.
+    """
+    embed = HuggingFaceEmbeddings("sentence-transformers/all-mpnet-base-v2")
+    supa = create_client(
+        os.getenv("SUPABASE_URL"), os.getenv("SUPABASE_SERVICE_KEY")
+    )
+    store = SupabaseVectorStore(
         client=supa,
-        embedding=embed_model,
+        embedding=embed,
         table_name="documents",
-        query_name="match_documents_langchain"
+        query_name="match_documents_langchain",
     )
-    return vs.as_retriever()
+    return store.as_retriever()
+
+
+def get_llm(provider: str = "google"):
+    """
+    Factory to select and return an LLM client.
+
+    Args:
+        provider: One of "google", "groq", "huggingface".
 
-# --- LLM Factory ---
-def get_llm(provider: str):
+    Returns:
+        Configured LLM client.
+
+    Raises:
+        ValueError: On unsupported provider.
+    """
     if provider == "google":
-        return ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+        return ChatGoogleGenerativeAI("gemini-2.0-flash", temperature=0)
     if provider == "groq":
-        return ChatGroq(model="qwen-qwq-32b", temperature=0)
+        return ChatGroq("qwen-qwq-32b", temperature=0)
     if provider == "huggingface":
-        return ChatHuggingFace(llm=HuggingFaceEndpoint(
-            url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
-            temperature=0))
+        return ChatHuggingFace(
+            llm=HuggingFaceEndpoint(
+                url="https://api-inference.huggingface.co/models/"
+                "Meta-DeepLearning/llama-2-7b-chat-hf",
+                temperature=0,
+            )
+        )
     raise ValueError(f"Unsupported provider: {provider}")
 
-# --- Build Graph ---
+
 def build_graph(provider: str = "google"):
-    # tools list
+    """
+    Build and compile a StateGraph for retrieval + LLM responses.
+
+    Args:
+        provider: LLM provider key.
+
+    Returns:
+        A compiled StateGraph.
+    """
+    sys_msg = load_system_prompt()
     retriever = build_vector_retriever()
     question_tool = create_retriever_tool(
         retriever=retriever,
         name="Question Search",
-        description="Retrieve similar Q&A from vector store"
+        description="Retrieve similar Q&A from vector store.",
     )
+
     tools = [
-        add, subtract, multiply, divide, modulus,
-        wiki_search, web_search, arvix_search,
-        question_tool
+        add,
+        subtract,
+        multiply,
+        divide,
+        modulus,
+        wiki_search,
+        web_search,
+        arxiv_search,
+        question_tool,
     ]
-
-    # LLM w/ tools
     llm = get_llm(provider).bind_tools(tools)
 
-    # Nodes
-    def assistant(state: MessagesState):
-        msgs = [sys_msg] + state["messages"]
-        resp = llm.invoke({"messages": msgs})
-        return {"messages": [resp]}
-
-    def retriever_node(state: MessagesState):
+    def retriever_node(state: MessagesState) -> dict:
+        """
+        Node: retrieve most relevant doc and extract its answer.
+        """
         query = state["messages"][-1].content
         doc = retriever.similarity_search(query, k=1)[0]
         text = doc.page_content
-        answer = text.split("Final answer :")[-1].strip() if "Final answer :" in text else text
-        return {"messages": [AIMessage(content=answer)]}
+        ans = text.split("Final answer :")[-1].strip() if "Final answer :" in text else text
+        return {"messages": [AIMessage(content=ans)]}
+
+    def assistant_node(state: MessagesState) -> dict:
+        """
+        Node: call LLM with system prompt + history.
+        """
+        msgs = [sys_msg] + state["messages"]
+        resp = llm.invoke({"messages": msgs})
+        return {"messages": [resp]}
 
-    # Graph assembly
     graph = StateGraph(MessagesState)
     graph.add_node("retriever", retriever_node)
-    graph.add_node("assistant", assistant)
+    graph.add_node("assistant", assistant_node)
     graph.add_node("tools", ToolNode(tools))
     graph.add_edge(START, "retriever")
     graph.add_edge("retriever", "assistant")
@@ -135,4 +239,4 @@ def build_graph(provider: str = "google"):
     graph.set_entry_point("retriever")
     graph.set_finish_point("assistant")
 
-    return graph.compile()
+    return graph.compile()