agent.py

"""
LangGraph Agent - 多工具智能代理系统
结合数学计算、网络搜索、学术检索和向量数据库增强能力
支持多种AI模型提供商（Google Gemini, Groq, HuggingFace）
"""

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
from langchain_core.tools import tool
from langchain.tools.retriever import create_retriever_tool
from supabase.client import Client, create_client

# 加载环境变量（API密钥、数据库连接等）
load_dotenv()

# ======================
# 工具定义部分 加减乘除搜索等
# ======================

@tool
def multiply(a: int, b: int) -> int:
    """乘法运算: 返回两个整数的乘积"""
    return a * b

@tool
def add(a: int, b: int) -> int:
    """加法运算: 返回两个整数的和"""
    return a + b

@tool
def subtract(a: int, b: int) -> int:
    """减法运算: 返回两个整数的差"""
    return a - b

@tool
def divide(a: int, b: int) -> int:
    """除法运算: 返回两个整数的商"""
    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:
    """维基百科搜索: 返回最多1个相关结果"""
    search_docs = WikipediaLoader(query=query, load_max_docs=1).load()
    # 格式化搜索结果
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ])
    return {"wiki_results": formatted_search_docs}

@tool
def web_search(query: str) -> str:
    """网络搜索(Tavily): 返回最多1个相关结果"""
    search_docs = TavilySearchResults(max_results=1).invoke(query=query)
    # 格式化搜索结果
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ])
    return {"web_results": formatted_search_docs}

@tool
def arvix_search(query: str) -> str:
    """学术论文搜索(Arxiv): 返回最多1个相关结果"""
    search_docs = ArxivLoader(query=query, load_max_docs=1).load()
    # 格式化搜索结果（截取前1000字符）
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
            for doc in search_docs
        ])
    return {"arvix_results": formatted_search_docs}

# ======================
# 系统初始化和配置
# ======================

# 加载系统提示（定义AI行为准则）
with open("system_prompt.txt", "r", encoding="utf-8") as f:
    system_prompt = f.read()
sys_msg = SystemMessage(content=system_prompt)

# 构建向量数据库检索工具
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
supabase: Client = 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",
)
retriever_tool = create_retriever_tool(
    retriever=vector_store.as_retriever(),
    name="Question Search",
    description="从向量数据库中检索相似问题",
)

# 所有可用工具列表
tools = [
    multiply,
    add,
    subtract,
    divide,
    modulus,
    wiki_search,
    web_search,
    arvix_search,
    retriever_tool,  # 添加向量检索工具
]

# ======================
# 图构建函数（核心逻辑）
# ======================

def build_graph(provider: str = "groq"):
    """
    构建LangGraph工作流
    
    参数:
        provider: AI模型提供商 ("google", "groq", "huggingface")
    
    返回:
        编译好的LangGraph对象
    """
    
    # 1. 选择AI模型
    if provider == "google":
        # Google Gemini模型
        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
    elif provider == "groq":
        # Groq高速推理模型
        llm = ChatGroq(model="qwen-qwq-32b", temperature=0)  # 可选模型: qwen-qwq-32b, gemma2-9b-it
    elif provider == "huggingface":
        # HuggingFace端点模型
        llm = ChatHuggingFace(
            llm=HuggingFaceEndpoint(
                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
                temperature=0,
            ),
        )
    else:
        raise ValueError("无效的提供商。请选择 'google', 'groq' 或 'huggingface'")
    
    # 2. 将工具绑定到AI模型
    llm_with_tools = llm.bind_tools(tools)

    # 3. 定义图节点
    
    def retriever_node(state: MessagesState):
        """检索节点：从向量数据库查找相似问题"""
        # 获取最新用户消息
        user_query = state["messages"][-1].content
        
        # 从向量数据库检索相似问题
        similar_question = vector_store.similarity_search(user_query)
        
        # 构建参考消息
        reference_msg = HumanMessage(
            content=f"参考类似问题及解答:\n\n{similar_question[0].page_content}",
        )
        
        # 返回增强后的消息流：系统提示 + 原始消息 + 参考消息
        return {"messages": [sys_msg] + state["messages"] + [reference_msg]}

    def assistant_node(state: MessagesState):
        """AI节点：处理消息并决定下一步动作"""
        # 调用AI模型处理当前消息状态
        response = llm_with_tools.invoke(state["messages"])
        return {"messages": [response]}

    # 4. 构建图结构
    builder = StateGraph(MessagesState)
    
    # 添加节点
    builder.add_node("retriever", retriever_node)  # 检索节点
    builder.add_node("assistant", assistant_node)  # AI处理节点
    builder.add_node("tools", ToolNode(tools))     # 工具执行节点
    
    # 设置节点间关系
    builder.add_edge(START, "retriever")  # 开始 -> 检索
    builder.add_edge("retriever", "assistant")  # 检索 -> AI处理
    
    # 条件边：AI处理后判断是否需要调用工具
    builder.add_conditional_edges(
        "assistant",
        tools_condition,  # 内置工具判断条件
    )
    
    # 工具执行后返回AI节点
    builder.add_edge("tools", "assistant")
    
    # 5. 编译图
    return builder.compile()

# ======================
# 测试执行
# ======================

if __name__ == "__main__":
    # 测试问题
    question = "托马斯·阿奎纳斯的图片是什么时候首次添加到双重效应原则的维基百科页面的？"
    
    # 构建图（使用Groq提供商）
    agent_graph = build_graph(provider="groq")
    
    # 初始化消息
    messages = [HumanMessage(content=question)]
    
    # 执行图工作流
    result = agent_graph.invoke({"messages": messages})
    
    # 打印所有消息
    print("\n===== 完整对话记录 =====")
    for msg in result["messages"]:
        print(f"[{msg.type}]: {msg.content[:200]}...")
    
    # 提取最终回答
    final_answer = result["messages"][-1].content
    print("\n===== 最终回答 =====")
    print(final_answer)