Update agent.py

agent.py

@@ -1,71 +1,123 @@
 # agent.py
 import os
 from dotenv import load_dotenv
-from langgraph.graph import StateGraph, MessagesState
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import tools_condition, ToolNode
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_community.vectorstores import SupabaseVectorStore
 from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langchain_core.tools import tool
+from langchain.tools.retriever import create_retriever_tool
+from supabase.client import create_client
 
-# Load environment variables (API keys, etc.)
 load_dotenv()
 
-# === TOOLS ===
+# ========== TOOLS ==========
+@tool
+def multiply(a: int, b: int) -> int:
+    return a * b
 
 @tool
 def add(a: int, b: int) -> int:
-    """Add two numbers."""
     return a + b
 
+@tool
+def subtract(a: int, b: int) -> int:
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> float:
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    return a % b
+
 @tool
 def wiki_search(query: str) -> str:
-    """Search Wikipedia and return up to 2 results."""
     search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    formatted = "\n\n---\n\n".join(
-        [f"{doc.metadata}\n{doc.page_content[:500]}" for doc in search_docs]
-    )
-    return {"wiki_results": formatted}
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'{doc.metadata}\n{doc.page_content[:500]}'
+            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 max 3 results."""
     search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-    formatted = "\n\n---\n\n".join(
-        [f"{doc.metadata}\n{doc.page_content[:500]}" for doc in search_docs]
-    )
-    return {"web_results": formatted}
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'{doc.metadata}\n{doc.page_content[:500]}'
+            for doc in search_docs
+        ])
+    return {"web_results": formatted_search_docs}
 
-tools = [add, wiki_search, web_search]
+@tool
+def arvix_search(query: str) -> str:
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted_search_docs = "\n\n---\n\n".join(
+        [
+            f'{doc.metadata}\n{doc.page_content[:1000]}'
+            for doc in search_docs
+        ])
+    return {"arvix_results": formatted_search_docs}
+
+tools = [
+    multiply, add, subtract, divide, modulus,
+    wiki_search, web_search, arvix_search,
+]
 
-# === SYSTEM PROMPT ===
+# ========== SYSTEM PROMPT ==========
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
     system_prompt = f.read()
 sys_msg = SystemMessage(content=system_prompt)
 
-# === GRAPH BUILD ===
-def build_graph():
-    # For demo, use OpenAI API if key provided; else default to Gemini
-    from langchain_openai import ChatOpenAI
-    openai_api_key = os.environ.get("OPENAI_API_KEY")
-    if openai_api_key:
-        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, openai_api_key=openai_api_key)
-    else:
-        # Gemini fallback (or any other, e.g., groq, etc.)
-        from langchain_google_genai import ChatGoogleGenerativeAI
-        llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
+# ========== SUPABASE VECTORSTORE ==========
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+supabase = create_client(
+    os.environ.get("SUPABASE_URL"),
+    os.environ.get("SUPABASE_SERVICE_KEY")
+)
+vector_store = SupabaseVectorStore(
+    client=supabase,
+    embedding=embeddings,
+    table_name="documents",
+    query_name="match_documents_langchain",
+)
 
-    # Bind tools to LLM
+# ========== GRAPH BUILD ==========
+def build_graph(provider: str = "google"):
+    if provider == "google":
+        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+    else:
+        raise ValueError("Invalid provider. (Expand for openai/groq if needed)")
     llm_with_tools = llm.bind_tools(tools)
 
-    # Simple assistant node
+    def retriever(state: MessagesState):
+        query = state["messages"][-1].content
+        similar_doc = vector_store.similarity_search(query, k=1)[0]
+        content = similar_doc.page_content
+        # Try to extract final answer format
+        if "FINAL ANSWER:" in content:
+            answer = content.split("FINAL ANSWER:")[-1].strip()
+            return {"messages": [AIMessage(content=f"FINAL ANSWER: {answer}")]}
+        else:
+            return {"messages": [AIMessage(content=content.strip())]}
+
     def assistant(state: MessagesState):
-        # Always include system prompt!
         messages = [sys_msg] + state["messages"]
         return {"messages": [llm_with_tools.invoke(messages)]}
 
     builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
-    builder.set_entry_point("assistant")
+    builder.add_edge(START, "retriever")
+    builder.add_edge("retriever", "assistant")
     builder.set_finish_point("assistant")
-
     return builder.compile()