update agent

agent.py

@@ -1,22 +1,25 @@
 import os
 from dotenv import load_dotenv
-from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.graph import START, StateGraph, MessagesState, END
 from langgraph.prebuilt import tools_condition, ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
-from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
 from langchain_core.tools import tool
 from langchain_groq import ChatGroq
+from supabase import create_client
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.vectorstores import SupabaseVectorStore
+from langchain_openai import ChatOpenAI
+from langchain_core.documents import Document
+import json
 
 load_dotenv()
 
-
-
 # ------------------- TOOL DEFINITIONS -------------------
-
 @tool
 def multiply(a: int, b: int) -> int:
     """Multiply two numbers."""
@@ -29,19 +32,19 @@ def add(a: int, b: int) -> int:
 
 @tool
 def subtract(a: int, b: int) -> int:
-    """Subtract two numbers."""
+    """Subtract b from a."""
     return a - b
 
 @tool
 def divide(a: int, b: int) -> float:
-    """Divide two numbers."""
+    """Divide a by b. Raise error if b is zero."""
     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."""
+    """Get remainder of a divided by b."""
     return a % b
 
 @tool
@@ -54,23 +57,19 @@ def wiki_search(query: str) -> str:
 def web_search(query: str) -> str:
     """Search the web using Tavily (max 3 results)."""
     results = TavilySearchResults(max_results=3).invoke(query)
-    texts = []
-    for doc in results:
-        if isinstance(doc, dict):
-            texts.append(doc.get("content", "") or doc.get("text", ""))
+    texts = [doc.get("content", "") or doc.get("text", "") for doc in results if isinstance(doc, dict)]
     return "\n\n".join(texts)
 
-
 @tool
 def arvix_search(query: str) -> str:
-    """Search Arxiv for academic papers (max 3)."""
+    """Search Arxiv for academic papers (max 3 results, truncated to 1000 characters each)."""
     docs = ArxivLoader(query=query, load_max_docs=3).load()
     return "\n\n".join([doc.page_content[:1000] for doc in docs])
 
+
 tools = [multiply, add, subtract, divide, modulus, wiki_search, web_search, arvix_search]
 
 # ------------------- SYSTEM PROMPT -------------------
-
 system_prompt_path = "system_prompt.txt"
 if os.path.exists(system_prompt_path):
     with open(system_prompt_path, "r", encoding="utf-8") as f:
@@ -83,12 +82,7 @@ else:
 sys_msg = SystemMessage(content=system_prompt)
 
 # ------------------- GRAPH CONSTRUCTION -------------------
-
-from langchain_openai import ChatOpenAI  # ✅ 新增导入
-
 def build_graph(provider: str = "groq"):
-    """Build LangGraph agent with QA retriever and tool-use fallback."""
-    # 初始化 LLM
     if provider == "google":
         llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
     elif provider == "groq":
@@ -111,13 +105,11 @@ def build_graph(provider: str = "groq"):
     else:
         raise ValueError("Invalid provider")
 
-    # 工具绑定
     llm_with_tools = llm.bind_tools(tools)
 
     def assistant(state: MessagesState):
         return {"messages": [sys_msg] + [llm_with_tools.invoke(state["messages"])]}
 
-    # ✅ 初始化 Supabase Retriever
     SUPABASE_URL = os.getenv("SUPABASE_URL")
     SUPABASE_KEY = os.getenv("SUPABASE_SERVICE_KEY")
     supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
@@ -130,7 +122,38 @@ def build_graph(provider: str = "groq"):
     )
     retriever = vectorstore.as_retriever(search_kwargs={"k": 1})
 
-    # ✅ Retriever 节点
+
+    # ✅ 替换 similarity_search_by_vector_with_relevance_scores 方法，直接调用 supabase.rpc
+    original_fn = vectorstore.similarity_search_by_vector_with_relevance_scores
+
+    # ✅ 覆盖 vectorstore 的方法
+    def patched_fn(embedding, k=4, filter=None, **kwargs):
+        response = supabase.rpc(
+            "match_documents",
+            {
+                "query_embedding": embedding,
+                "match_count": k
+            }
+        ).execute()
+
+        documents = []
+        for r in response.data:
+            metadata = r["metadata"]
+            if isinstance(metadata, str):
+                try:
+                    metadata = json.loads(metadata)
+                except Exception:
+                    metadata = {}
+            doc = Document(
+                page_content=r["content"],
+                metadata=metadata
+            )
+            documents.append((doc, r["similarity"]))
+        return documents
+
+    # ✅ 覆盖 vectorstore 的方法
+    vectorstore.similarity_search_by_vector_with_relevance_scores = patched_fn
+
     def qa_retriever_node(state: MessagesState):
         user_question = state["messages"][-1].content
         docs = retriever.invoke(user_question)
@@ -139,12 +162,8 @@ def build_graph(provider: str = "groq"):
                 "messages": state["messages"] + [AIMessage(content=docs[0].page_content)],
                 "__condition__": "complete"
             }
-        return {
-            "messages": state["messages"],
-            "__condition__": "default"
-        }
+        return {"messages": state["messages"], "__condition__": "default"}
 
-    # 构建图结构
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", qa_retriever_node)
     builder.add_node("assistant", assistant)
@@ -152,8 +171,8 @@ def build_graph(provider: str = "groq"):
 
     builder.add_edge(START, "retriever")
     builder.add_conditional_edges("retriever", {
-        "default": "assistant",
-        "complete": None
+        "default": lambda x: "assistant",
+        "complete": lambda x: END,
     })
     builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")
@@ -161,7 +180,6 @@ def build_graph(provider: str = "groq"):
     return builder.compile()
 
 # ------------------- LOCAL 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?"
     graph = build_graph(provider="openai")