add retriever

agent.py

@@ -87,7 +87,8 @@ sys_msg = SystemMessage(content=system_prompt)
 from langchain_openai import ChatOpenAI  # ✅ 新增导入
 
 def build_graph(provider: str = "groq"):
-    """Build the LangGraph with tool-use."""
+    """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":
@@ -106,21 +107,54 @@ def build_graph(provider: str = "groq"):
         openai_key = os.getenv("OPENAI_API_KEY")
         if not openai_key:
             raise ValueError("OPENAI_API_KEY is not set.")
-        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, api_key=openai_key)  # ✅ OpenAI GPT
+        llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, api_key=openai_key)
     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"])]}
 
-    # Build the graph with assistant and tools
+    # ✅ 初始化 Supabase Retriever
+    SUPABASE_URL = os.getenv("SUPABASE_URL")
+    SUPABASE_KEY = os.getenv("SUPABASE_SERVICE_KEY")
+    supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
+
+    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+    vectorstore = SupabaseVectorStore(
+        client=supabase,
+        embedding=embedding_model,
+        table_name="QA_db"
+    )
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 1})
+
+    # ✅ Retriever 节点
+    def qa_retriever_node(state: MessagesState):
+        user_question = state["messages"][-1].content
+        docs = retriever.invoke(user_question)
+        if docs:
+            return {
+                "messages": state["messages"] + [AIMessage(content=docs[0].page_content)],
+                "__condition__": "complete"
+            }
+        return {
+            "messages": state["messages"],
+            "__condition__": "default"
+        }
+
+    # 构建图结构
     builder = StateGraph(MessagesState)
+    builder.add_node("retriever", qa_retriever_node)
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(tools))
 
-    builder.add_edge(START, "assistant")
+    builder.add_edge(START, "retriever")
+    builder.add_conditional_edges("retriever", {
+        "default": "assistant",
+        "complete": None
+    })
     builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")
 

requirements.txt

@@ -17,3 +17,4 @@ pymupdf
 wikipedia
 pgvector
 python-dotenv
+supabase