"""LangGraph Agent""" 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 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.vectorstores import SupabaseVectorStore from langchain_core.messages import SystemMessage, HumanMessage, AIMessage from langchain_core.tools import tool from langchain_core.vectorstores import InMemoryVectorStore from langchain.tools.retriever import create_retriever_tool load_dotenv() @tool def multiply(a: int, b: int) -> int: """Multiply two numbers. Args: a: first int b: second int """ return a * b @tool def add(a: int, b: int) -> int: """Add two numbers. Args: a: first int b: second int """ return a + b @tool def subtract(a: int, b: int) -> int: """Subtract two numbers. Args: a: first int b: second int """ return a - b @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 @tool def modulus(a: int, b: int) -> int: """Get the modulus of two numbers. Args: a: first int b: second int """ return a % b @tool def wiki_search(query: str) -> str: """Search Wikipedia for a query and return maximum 2 results. Args: query: The search query.""" search_docs = WikipediaLoader(query=query, load_max_docs=2).load() formatted_search_docs = "\n\n---\n\n".join( [ f'\n{doc.page_content}\n' for doc in search_docs ]) return {"wiki_results": formatted_search_docs} @tool def web_search(query: str) -> str: """Search Tavily for a query and return maximum 3 results. Args: query: The search query.""" search_docs = TavilySearchResults(max_results=3).invoke(query=query) formatted_search_docs = "\n\n---\n\n".join( [ f'\n{doc.page_content}\n' for doc in search_docs ]) return {"web_results": formatted_search_docs} @tool def arvix_search(query: str) -> str: """Search Arxiv for a query and return maximum 3 result. Args: query: The search query.""" search_docs = ArxivLoader(query=query, load_max_docs=3).load() formatted_search_docs = "\n\n---\n\n".join( [ f'\n{doc.page_content[:1000]}\n' for doc in search_docs ]) return {"arvix_results": formatted_search_docs} # load the system prompt from the file with open("system_prompt.txt", "r", encoding="utf-8") as f: system_prompt = f.read() # System message sys_msg = SystemMessage(content=system_prompt) # build a 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")) vector_store = InMemoryVectorStore(embeddings) # Add sample documents to the vector store from langchain_core.documents import Document import pandas as pd import ast with open("supabase_docs.csv", "r", encoding="utf-8") as f: df = pd.read_csv(f) documents = [] for _, row in df.iterrows(): content = row["content"] # parse the metadata string into a dict metadata = ast.literal_eval(row["metadata"]) documents.append(Document(page_content=content, metadata=metadata)) vector_store.add_documents(documents) create_retriever_tool = create_retriever_tool( retriever=vector_store.as_retriever(), name="question_search", description="A tool to retrieve similar questions from a vector store.", ) tools = [ multiply, add, subtract, divide, modulus, wiki_search, web_search, arvix_search, create_retriever_tool, ] # Build graph function def build_graph(provider: str = "google"): """Build the graph""" # Load environment variables from .env file if provider == "google": # Google Gemini llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0) elif provider == "groq": # Groq https://console.groq.com/docs/models llm = ChatGroq(model="qwen-qwq-32b", temperature=0) # optional : qwen-qwq-32b gemma2-9b-it else: raise ValueError("Invalid provider. Choose 'google' or 'groq'.") # Bind tools to LLM llm_with_tools = llm.bind_tools(tools) # Node def assistant(state: MessagesState): last = state["messages"][-1] # If retriever has already answered, do nothing. if isinstance(last, AIMessage) and last.content.startswith("FINAL ANSWER"): return {"messages": state["messages"]} # short‑circuit # Otherwise call the LLM as usual return {"messages": [llm_with_tools.invoke(state["messages"])]} def retriever(state: MessagesState): query = state["messages"][-1].content hits = vector_store.similarity_search(query, k=1) if not hits: return {"messages": state["messages"]} example = hits[0].page_content.strip() # your stored Q&A blob # Split into question / final answer q_part, a_part = example.split("Final answer :") demo_q = HumanMessage(content=q_part.strip()) demo_ans = AIMessage(content=f"FINAL ANSWER: {a_part.strip()}") # Only prepend sys_msg once base = [] if isinstance(state["messages"][0], SystemMessage) else [sys_msg] # **Order matters** – give few‑shot demo *before* the real question new_messages = base + [demo_q, demo_ans] + state["messages"] return {"messages": new_messages} builder = StateGraph(MessagesState) builder.add_node("retriever", retriever) builder.add_node("assistant", assistant) builder.add_node("tools", ToolNode(tools)) builder.add_edge(START, "retriever") builder.add_edge("retriever", "assistant") builder.add_conditional_edges( "assistant", tools_condition, ) builder.add_edge("tools", "assistant") # Compile graph return builder.compile() # test if __name__ == "__main__": question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?" # Build the graph graph = build_graph(provider="google") # Run the graph messages = [HumanMessage(content=question)] messages = graph.invoke({"messages": messages}) for m in messages["messages"]: m.pretty_print()