Update agent.py

agent.py

@@ -1,171 +1,160 @@
 # agent.py
+
 import os
-#from supabase import create_client
-from sentence_transformers import SentenceTransformer
-from serpapi import GoogleSearch
 import pandas as pd
 import faiss
-from langgraph.graph import Graph
+
+from sentence_transformers import SentenceTransformer
+from serpapi import GoogleSearch
+
+# 1️⃣ Switch Graph → StateGraph
+from langgraph.graph import StateGraph
 from langchain_core.language_models.llms import LLM
-from langchain_core.tools import tool 
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import tool
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
 
-# ─── 1) Load & embed all documents at startup ───
-# 1a) Read CSV of docs
+# ────────────────
+# 2️⃣ Load & index your static FAISS docs
+# ────────────────
 df = pd.read_csv("documents.csv")
 DOCS = df["content"].tolist()
 
-# 1b) Create an embedding model
 EMBEDDER = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")
-
-# 1c) Compute embeddings (float32) and build FAISS index
 EMBS = EMBEDDER.encode(DOCS, show_progress_bar=True).astype("float32")
 INDEX = faiss.IndexFlatL2(EMBS.shape[1])
 INDEX.add(EMBS)
 
+# ────────────────
+# 3️⃣ Read your system prompt
+# ────────────────
+with open("system_prompt.txt","r",encoding="utf-8") as f:
+    SYSTEM_PROMPT = f.read().strip()
 
-# ----Supabase setup----
-SUPABASE_URL = os.getenv("SUPABASE_URL")
-SUPABASE_KEY = os.getenv("SUPABASE_SERVICE_KEY")
-EMBED_MODEL_ID = os.getenv("HF_EMBEDDING_MODEL")
-
-
-#sb_client = create_client(SUPABASE_URL, SUPABASE_KEY)
-#embedder = SentenceTransformer(EMBED_MODEL_ID)
-
-# 1) Define tools
-
+# ────────────────
+# 4️⃣ Define your tools (unchanged semantics)
+# ────────────────
 @tool
 def calculator(expr: str) -> str:
-    """Simple math via Python eval"""
     try:
         return str(eval(expr))
-    except Exception:
+    except:
         return "Error"
-# e.g. search, vector_retrieval, etc.
-# @tool
-# def web_search(query:str) -> str:
-# ...
-#@tool
-#def retrieve_docs(query: str, k: int = 3) -> str:
-    #"""
-    #Fetch tpo-k docs from Supabase vector store.
-    #Returns the concatenated text.
-    #"""
-    # --- embed the query
-    #q_emb = embedder.encode(query).tolist()
-
-    # --- query the embedding table
-    #response = (
-   #     sb_client
-   #     .rpc("match_documents", {"query_embedding": q_emb, "match_count": k})
-   #     .execute()
-   # )
-   # rows = response.data
-
-    # ---- concatenate the content field
-   # docs = [row["content"] for row in rows]
-   # return "\n\n---\n\n".join(docs)
 
 @tool
 def retrieve_docs(query: str, k: int = 3) -> str:
-    """
-    k-NN search over our in-memory FAISS index.
-    Returns the top-k documents concatenated.
-    """
-    # 1) Embed the query
     q_emb = EMBEDDER.encode([query]).astype("float32")
-    # 2) Search FAISS
     D, I = INDEX.search(q_emb, k)
-    # 3) Gather and return the texts
-    hits = [DOCS[i] for i in I[0]]
-    return "\n\n---\n\n".join(hits)
+    return "\n\n---\n\n".join(DOCS[i] for i in I[0])
 
-    
 SERPAPI_KEY = os.getenv("SERPAPI_KEY")
-# ---- web_search tool
 @tool
 def web_search(query: str, num_results: int = 5) -> str:
-    """ Return top-5 snippets from Google search via SerpAPI."""
-    params = {
-        "engine": "google",
-        "q": query,
-        "num": num_results,
-        "api_key": SERPAPI_KEY,
-    }
-    search = GoogleSearch(params)
-    results = search.get_dict().get("organic_results", [])
-    snippets = [r.get("snippet","")for r in results]
-    return "\n".join(f"- {s}" for s in snippets)
+    params = {"engine":"google","q":query,"num":num_results,"api_key":SERPAPI_KEY}
+    res = GoogleSearch(params).get_dict().get("organic_results", [])
+    return "\n".join(f"- {r.get('snippet','')}" for r in res)
 
 @tool
 def wiki_search(query: str) -> str:
-    """
-    Search Wikipedia for up to 2 pages matching 'query',
-    and return their contents.
-    """
-    #load up to 2 pages
     pages = WikipediaLoader(query=query, load_max_docs=2).load()
-    #format as plain text
-    return "\n\n---\n\n".join(doc.page_content for doc in pages)
+    return "\n\n---\n\n".join(d.page_content for d in pages)
 
 @tool
-def arxiv_search(query:str) -> str:
-    """
-    Search ArXiv for up to 3 abstracts matching 'query',
-    and return their first 1000 characters.
-    """
+def arxiv_search(query: str) -> str:
     papers = ArxivLoader(query=query, load_max_docs=3).load()
-    return "\n\n---\n\n".join(doc.page_content[:1000]for doc in papers)
-
-
-#read the system prompt
-with open("system_prompt.txt","r",encoding="utf-8") as f:
-    SYSTEM_PROMPT = f.read()
-
-# 2) Build your graph
-def build_graph(provider: str = "huggingface") -> Graph:
-    # 2a) Instantiate your LLM endpoint
-    api_token = os.getenv("HF_TOKEN")
-        if not api_token:
-            raise ValueError("HF_TOKEN not found: please add it under Settings → Secrets and variables in your Space")
-    llm = LLM(provider=provider, token=api_token,  model="meta-llama/Llama-2-7b-chat-hf")
-
-    # 2b) Attach tools   
-    tools = [
-        calculator,
-        retrieve_docs,
-        web_search,
-        wiki_search,
-        arxiv_search, # add more tools here
-    ]
-    llm_with_tools = llm.bind_tools(tools) 
-
-    # 2c) Compose your graph
-    graph = Graph()
-    # sys node: prepend system prompt
-    def _prepend_system(query: str) -> str:
-        return SYSTEM_PROMPT + "\n\n" + query
-    graph.add_node("sys", _prepend_system)
-    # "ask" node : the LLM itself
-    graph.add_node("ask", llm_with_tools) # prompt node
-    # tool nodes
-    graph.add_node("calc", calculator)
-    graph.add_node("retrieve", retrieve_docs)
-    graph.add_node("web_search", web_search)
-    graph.add_node("wiki", wiki_search)
-    graph.add_node("arxiv", arxiv_search)
-    # allow the LLM to call any tool:
-    graph.add_edge("ask", "calc") # allow ask -> calc
-    graph.add_edge("ask", "retrieve")
-    graph.add_edge("ask", "web_search")
-    graph.add_edge("ask", "wiki")
-    graph.add_edge("ask", "arxiv")
-
-    # wire up the start: sys -> ask
-    graph.add_edge(Graph.START, "sys")
-    graph.add_edge("sys", "ask")
-    graph.set_start("ask")
-
-    return graph
-    
+    return "\n\n---\n\n".join(d.page_content[:1000] for d in papers)
+
+
+# ────────────────
+# 5️⃣ Define your State schema
+# ────────────────
+from typing import TypedDict, List
+from langchain_core.messages import BaseMessage
+
+class AgentState(TypedDict):
+    # We’ll carry a list of messages as our “chat history”
+    messages: List[BaseMessage]
+
+
+# ────────────────
+# 6️⃣ Build the StateGraph
+# ────────────────
+def build_graph(provider: str = "huggingface") -> StateGraph:
+    # Instantiate LLM
+    hf_token = os.getenv("HF_TOKEN")
+    if not hf_token:
+        raise ValueError("HF_TOKEN missing in env")
+    llm = LLM(provider=provider, token=hf_token, model="meta-llama/Llama-2-7b-chat-hf")
+
+    # 6.1) Node: init → seed system prompt
+    def init_node(_: AgentState) -> AgentState:
+        return {
+            "messages": [
+                SystemMessage(content=SYSTEM_PROMPT)
+            ]
+        }
+    
+    # 6.2) Node: human → append user question
+    def human_node(state: AgentState, question: str) -> AgentState:
+        state["messages"].append(HumanMessage(content=question))
+        return state
+
+    # 6.3) Node: assistant → call LLM on current messages
+    def assistant_node(state: AgentState) -> dict:
+        ai_msg = llm.invoke(state["messages"])
+        return {"messages": state["messages"] + [ai_msg]}
+
+    # 6.4) Optional: tool nodes (they’ll read last HumanMessage)
+    def make_tool_node(fn):
+        def tool_node(state: AgentState) -> dict:
+            # fetch the latest human query
+            last_query = state["messages"][-1].content
+            result = fn(last_query)
+            # append the tool’s output as if from system/Human
+            state["messages"].append(HumanMessage(content=result))
+            return {"messages": state["messages"]}
+        return tool_node
+
+    # Instantiate nodes for each tool
+    calc_node     = make_tool_node(calculator)
+    retrieve_node = make_tool_node(retrieve_docs)
+    web_node      = make_tool_node(web_search)
+    wiki_node     = make_tool_node(wiki_search)
+    arxiv_node    = make_tool_node(arxiv_search)
+
+    # 6.5) Build the graph
+    g = StateGraph(AgentState)
+
+    # Register nodes
+    g.add_node("init",    init_node)
+    g.add_node("human",   human_node)
+    g.add_node("assistant", assistant_node)
+    g.add_node("calc",    calc_node)
+    g.add_node("retrieve", retrieve_node)
+    g.add_node("web",      web_node)
+    g.add_node("wiki",     wiki_node)
+    g.add_node("arxiv",    arxiv_node)
+
+    # Wire up edges
+    from langgraph.graph import END
+    g.set_entry_point("init")
+    # init → human (placeholder: we’ll inject the actual question at runtime)
+    g.add_edge("init", "human")
+    # human → assistant
+    g.add_edge("human", "assistant")
+    # assistant → tool nodes (conditional on tool calls)
+    g.add_edge("assistant", "calc")
+    g.add_edge("assistant", "retrieve")
+    g.add_edge("assistant", "web")
+    g.add_edge("assistant", "wiki")
+    g.add_edge("assistant", "arxiv")
+    # each tool returns back into assistant for follow‐up
+    g.add_edge("calc",     "assistant")
+    g.add_edge("retrieve", "assistant")
+    g.add_edge("web",      "assistant")
+    g.add_edge("wiki",     "assistant")
+    g.add_edge("arxiv",    "assistant")
+    # and finally assistant → END when done
+    g.add_edge("assistant", END)
+
+    return g.compile()