Update agent.py

agent.py

@@ -1,17 +1,14 @@
 import os
+import functools
 from dotenv import load_dotenv
-from supabase.client import create_client
+from supabase.client import create_client, Client
 from langgraph.graph import START, StateGraph, MessagesState
-from langgraph.prebuilt import ToolNode, tools_condition
+from langgraph.prebuilt import tools_condition, ToolNode
 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
@@ -19,224 +16,173 @@ from langchain.tools.retriever import create_retriever_tool
 
 load_dotenv()
 
+def _format_search_results(docs, label: str, truncate: int = None) -> dict:
+    """Helper to format document search results."""
+    entries = []
+    for d in docs:
+        content = d.page_content if truncate is None else d.page_content[:truncate]
+        entries.append(
+            f'<Document source="{d.metadata.get("source","")}" '
+            f'page="{d.metadata.get("page","")}"/>\n{content}\n</Document>'
+        )
+    return {label: "\n\n---\n\n".join(entries)}
 
-def load_system_prompt(path: str = "system_prompt.txt") -> SystemMessage:
-    """
-    Load system prompt from a file, fallback to a default if missing.
-
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers.
     Args:
-        path: File path to the system prompt.
-
-    Returns:
-        SystemMessage containing the loaded or default prompt.
+        a: first int
+        b: second int
     """
-    try:
-        with open(path, encoding="utf-8") as f:
-            content = f.read()
-    except FileNotFoundError:
-        content = "You are a helpful assistant."
-    return SystemMessage(content=content)
-
-
-def math_tool(func):
-    """
-    Wrap a Python function as a LangChain tool.
+    return a * b
 
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers.
+    
     Args:
-        func: Callable to wrap.
-
-    Returns:
-        A LangChain tool.
+        a: first int
+        b: second int
     """
-    return tool(func)
-
-
-@math_tool
-def add(a: int, b: int) -> int:
-    """Return a + b."""
     return a + b
 
-
-@math_tool
+@tool
 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."""
-    return a * b
-
-
-@math_tool
-def divide(a: int, b: int) -> float:
+    """Subtract two numbers.
+    
+    Args:
+        a: first int
+        b: second int
     """
-    Return a / b.
+    return a - b
 
-    Raises:
-        ValueError: If b is zero.
+@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
 
-
-@math_tool
+@tool
 def modulus(a: int, b: int) -> int:
-    """Return a % b."""
-    return a % b
-
-
-def format_docs(docs, key: str, max_chars: int = None) -> dict:
-    """
-    Convert document list into labeled XML-style chunks.
-
+    """Get the modulus of two numbers.
+    
     Args:
-        docs: Iterable of Document objects.
-        key: Dict key for formatted results.
-        max_chars: Optionally truncate content.
-
-    Returns:
-        {key: formatted_string}
+        a: first int
+        b: second int
     """
-    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)}
-
+    return a % b
 
 @tool
-def wiki_search(query: str) -> dict:
-    """Search Wikipedia (2 docs) and format results."""
+def wiki_search(query: str) -> str:
+    """Search Wikipedia for a query and return maximum 2 results.
+    
+    Args:
+        query: The search query."""
     docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    return format_docs(docs, "wiki_results")
-
+    return _format_search_results(docs, "wiki_results")
 
 @tool
-def web_search(query: str) -> dict:
-    """Search the web via Tavily (3 docs) and format results."""
+def web_search(query: str) -> str:
+    """Search Tavily for a query and return maximum 3 results.
+    
+    Args:
+        query: The search query."""
     docs = TavilySearchResults(max_results=3).invoke(query=query)
-    return format_docs(docs, "web_results")
-
+    return _format_search_results(docs, "web_results")
 
 @tool
-def arxiv_search(query: str) -> dict:
-    """Search ArXiv (3 docs) and format results (truncate to 1k chars)."""
+def arvix_search(query: str) -> str:
+    """Search Arxiv for a query and return maximum 3 result.
+    
+    Args:
+        query: The search query."""
     docs = ArxivLoader(query=query, load_max_docs=3).load()
-    return format_docs(docs, "arxiv_results", max_chars=1000)
-
-
-def build_vector_retriever():
-    """
-    Create and return a Supabase-based vector retriever.
+    return _format_search_results(docs, "arvix_results", truncate=1000)
 
-    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,
-        table_name="documents",
-        query_name="match_documents_langchain",
-    )
-    return store.as_retriever()
-
-
-def get_llm(provider: str = "google"):
-    """
-    Factory to select and return an LLM client.
+# load the system prompt from the file
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
 
-    Args:
-        provider: One of "google", "groq", "huggingface".
+# System message
+sys_msg = SystemMessage(content=system_prompt)
 
-    Returns:
-        Configured LLM client.
+# build a retriever once
+_embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+_supabase: Client = create_client(
+    os.environ["SUPABASE_URL"], os.environ["SUPABASE_SERVICE_KEY"]
+)
+_vector_store = SupabaseVectorStore(
+    client=_supabase,
+    embedding=_embeddings,
+    table_name="documents",
+    query_name="match_documents_langchain",
+)
+_retriever = _vector_store.as_retriever()
+_question_search_tool = create_retriever_tool(
+    retriever=_retriever,
+    name="Question Search",
+    description="A tool to retrieve similar questions from a vector store.",
+)
 
-    Raises:
-        ValueError: On unsupported provider.
-    """
-    if provider == "google":
-        return ChatGoogleGenerativeAI("gemini-2.0-flash", temperature=0)
-    if provider == "groq":
-        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,
-            )
+tools = [
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulus,
+    wiki_search,
+    web_search,
+    arvix_search,
+    _question_search_tool,
+]
+
+_LLM_PROVIDERS = {
+    "google": lambda: ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0),
+    "groq": lambda: ChatGroq(model="qwen-qwq-32b", temperature=0),
+    "huggingface": lambda: 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}")
-
+    ),
+}
 
-def build_graph(provider: str = "google"):
+@functools.lru_cache(maxsize=None)
+def get_llm(provider: str):
     """
-    Build and compile a StateGraph for retrieval + LLM responses.
-
-    Args:
-        provider: LLM provider key.
-
-    Returns:
-        A compiled StateGraph.
+    Retrieve and cache the LLM client for the given provider.
     """
-    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.",
-    )
-
-    tools = [
-        add,
-        subtract,
-        multiply,
-        divide,
-        modulus,
-        wiki_search,
-        web_search,
-        arxiv_search,
-        question_tool,
-    ]
+    try:
+        return _LLM_PROVIDERS[provider]()
+    except KeyError:
+        raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
+
+def build_graph(provider: str = "google"):
+    """Build the graph"""
     llm = get_llm(provider).bind_tools(tools)
 
-    def retriever_node(state: MessagesState) -> dict:
-        """
-        Node: retrieve most relevant doc and extract its answer.
-        """
+    def assistant(state: MessagesState):
+        """Assistant node"""
+        return {"messages": [llm.invoke(state["messages"])]}
+
+    def retriever(state: MessagesState):
         query = state["messages"][-1].content
-        doc = retriever.similarity_search(query, k=1)[0]
-        text = doc.page_content
-        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]}
+        doc = _retriever.similarity_search(query, k=1)[0]
+        content = doc.page_content
+        if "Final answer :" in content:
+            answer = content.split("Final answer :")[-1].strip()
+        else:
+            answer = content.strip()
+        return {"messages": [AIMessage(content=answer)]}
 
     graph = StateGraph(MessagesState)
-    graph.add_node("retriever", retriever_node)
-    graph.add_node("assistant", assistant_node)
-    graph.add_node("tools", ToolNode(tools))
-    graph.add_edge(START, "retriever")
-    graph.add_edge("retriever", "assistant")
-    graph.add_conditional_edges("assistant", tools_condition)
-    graph.add_edge("tools", "assistant")
+    graph.add_node("retriever", retriever)
     graph.set_entry_point("retriever")
-    graph.set_finish_point("assistant")
-
-    return graph.compile()
+    graph.set_finish_point("retriever")
+    return graph.compile()