|
|
"""
|
|
|
智能搜索助手 - 基于 LangGraph + Tavily API 的真实搜索系统
|
|
|
1. 理解用户需求
|
|
|
2. 使用Tavily API真实搜索信息
|
|
|
3. 生成基于搜索结果的回答
|
|
|
"""
|
|
|
|
|
|
import asyncio
|
|
|
from typing import TypedDict, Annotated
|
|
|
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
|
|
|
from langchain_openai import ChatOpenAI
|
|
|
from langgraph.graph import StateGraph, START, END
|
|
|
from langgraph.graph.message import add_messages
|
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
import os
|
|
|
from dotenv import load_dotenv
|
|
|
from tavily import TavilyClient
|
|
|
|
|
|
# 加载环境变量
|
|
|
load_dotenv()
|
|
|
|
|
|
# 定义状态结构
|
|
|
class SearchState(TypedDict):
|
|
|
messages: Annotated[list, add_messages]
|
|
|
user_query: str # 用户查询
|
|
|
search_query: str # 优化后的搜索查询
|
|
|
search_results: str # Tavily搜索结果
|
|
|
final_answer: str # 最终答案
|
|
|
step: str # 当前步骤
|
|
|
|
|
|
# 初始化模型和Tavily客户端
|
|
|
llm = ChatOpenAI(
|
|
|
model=os.getenv("LLM_MODEL_ID", "gpt-4o-mini"),
|
|
|
api_key=os.getenv("LLM_API_KEY"),
|
|
|
base_url=os.getenv("LLM_BASE_URL", "https://api.openai.com/v1"),
|
|
|
temperature=0.7
|
|
|
)
|
|
|
|
|
|
# 初始化Tavily客户端
|
|
|
tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
|
|
|
|
|
|
def understand_query_node(state: SearchState) -> SearchState:
|
|
|
"""步骤1：理解用户查询并生成搜索关键词"""
|
|
|
|
|
|
# 获取最新的用户消息
|
|
|
user_message = ""
|
|
|
for msg in reversed(state["messages"]):
|
|
|
if isinstance(msg, HumanMessage):
|
|
|
user_message = msg.content
|
|
|
break
|
|
|
|
|
|
understand_prompt = f"""分析用户的查询："{user_message}"
|
|
|
|
|
|
请完成两个任务：
|
|
|
1. 简洁总结用户想要了解什么
|
|
|
2. 生成最适合搜索的关键词（中英文均可，要精准）
|
|
|
|
|
|
格式：
|
|
|
理解：[用户需求总结]
|
|
|
搜索词：[最佳搜索关键词]"""
|
|
|
|
|
|
response = llm.invoke([SystemMessage(content=understand_prompt)])
|
|
|
|
|
|
# 提取搜索关键词
|
|
|
response_text = response.content
|
|
|
search_query = user_message # 默认使用原始查询
|
|
|
|
|
|
if "搜索词：" in response_text:
|
|
|
search_query = response_text.split("搜索词：")[1].strip()
|
|
|
elif "搜索关键词：" in response_text:
|
|
|
search_query = response_text.split("搜索关键词：")[1].strip()
|
|
|
|
|
|
return {
|
|
|
"user_query": response.content,
|
|
|
"search_query": search_query,
|
|
|
"step": "understood",
|
|
|
"messages": [AIMessage(content=f"我理解您的需求：{response.content}")]
|
|
|
}
|
|
|
|
|
|
def tavily_search_node(state: SearchState) -> SearchState:
|
|
|
"""步骤2：使用Tavily API进行真实搜索"""
|
|
|
|
|
|
search_query = state["search_query"]
|
|
|
|
|
|
try:
|
|
|
print(f"🔍 正在搜索: {search_query}")
|
|
|
|
|
|
# 调用Tavily搜索API
|
|
|
response = tavily_client.search(
|
|
|
query=search_query,
|
|
|
search_depth="basic",
|
|
|
include_answer=True,
|
|
|
include_raw_content=False,
|
|
|
max_results=5
|
|
|
)
|
|
|
|
|
|
# 处理搜索结果
|
|
|
search_results = ""
|
|
|
|
|
|
# 优先使用Tavily的综合答案
|
|
|
if response.get("answer"):
|
|
|
search_results = f"综合答案：\n{response['answer']}\n\n"
|
|
|
|
|
|
# 添加具体的搜索结果
|
|
|
if response.get("results"):
|
|
|
search_results += "相关信息：\n"
|
|
|
for i, result in enumerate(response["results"][:3], 1):
|
|
|
title = result.get("title", "")
|
|
|
content = result.get("content", "")
|
|
|
url = result.get("url", "")
|
|
|
search_results += f"{i}. {title}\n{content}\n来源：{url}\n\n"
|
|
|
|
|
|
if not search_results:
|
|
|
search_results = "抱歉，没有找到相关信息。"
|
|
|
|
|
|
return {
|
|
|
"search_results": search_results,
|
|
|
"step": "searched",
|
|
|
"messages": [AIMessage(content=f"✅ 搜索完成！找到了相关信息，正在为您整理答案...")]
|
|
|
}
|
|
|
|
|
|
except Exception as e:
|
|
|
error_msg = f"搜索时发生错误: {str(e)}"
|
|
|
print(f"❌ {error_msg}")
|
|
|
|
|
|
return {
|
|
|
"search_results": f"搜索失败：{error_msg}",
|
|
|
"step": "search_failed",
|
|
|
"messages": [AIMessage(content="❌ 搜索遇到问题，我将基于已有知识为您回答")]
|
|
|
}
|
|
|
|
|
|
def generate_answer_node(state: SearchState) -> SearchState:
|
|
|
"""步骤3：基于搜索结果生成最终答案"""
|
|
|
|
|
|
# 检查是否有搜索结果
|
|
|
if state["step"] == "search_failed":
|
|
|
# 如果搜索失败，基于LLM知识回答
|
|
|
fallback_prompt = f"""搜索API暂时不可用，请基于您的知识回答用户的问题：
|
|
|
|
|
|
用户问题：{state['user_query']}
|
|
|
|
|
|
请提供一个有用的回答，并说明这是基于已有知识的回答。"""
|
|
|
|
|
|
response = llm.invoke([SystemMessage(content=fallback_prompt)])
|
|
|
|
|
|
return {
|
|
|
"final_answer": response.content,
|
|
|
"step": "completed",
|
|
|
"messages": [AIMessage(content=response.content)]
|
|
|
}
|
|
|
|
|
|
# 基于搜索结果生成答案
|
|
|
answer_prompt = f"""基于以下搜索结果为用户提供完整、准确的答案：
|
|
|
|
|
|
用户问题：{state['user_query']}
|
|
|
|
|
|
搜索结果：
|
|
|
{state['search_results']}
|
|
|
|
|
|
请要求：
|
|
|
1. 综合搜索结果，提供准确、有用的回答
|
|
|
2. 如果是技术问题，提供具体的解决方案或代码
|
|
|
3. 引用重要信息的来源
|
|
|
4. 回答要结构清晰、易于理解
|
|
|
5. 如果搜索结果不够完整，请说明并提供补充建议"""
|
|
|
|
|
|
response = llm.invoke([SystemMessage(content=answer_prompt)])
|
|
|
|
|
|
return {
|
|
|
"final_answer": response.content,
|
|
|
"step": "completed",
|
|
|
"messages": [AIMessage(content=response.content)]
|
|
|
}
|
|
|
|
|
|
# 构建搜索工作流
|
|
|
def create_search_assistant():
|
|
|
workflow = StateGraph(SearchState)
|
|
|
|
|
|
# 添加三个节点
|
|
|
workflow.add_node("understand", understand_query_node)
|
|
|
workflow.add_node("search", tavily_search_node)
|
|
|
workflow.add_node("answer", generate_answer_node)
|
|
|
|
|
|
# 设置线性流程
|
|
|
workflow.add_edge(START, "understand")
|
|
|
workflow.add_edge("understand", "search")
|
|
|
workflow.add_edge("search", "answer")
|
|
|
workflow.add_edge("answer", END)
|
|
|
|
|
|
# 编译图
|
|
|
memory = InMemorySaver()
|
|
|
app = workflow.compile(checkpointer=memory)
|
|
|
|
|
|
return app
|
|
|
|
|
|
async def main():
|
|
|
"""主函数：运行智能搜索助手"""
|
|
|
|
|
|
# 检查API密钥
|
|
|
if not os.getenv("TAVILY_API_KEY"):
|
|
|
print("❌ 错误：请在.env文件中配置TAVILY_API_KEY")
|
|
|
return
|
|
|
|
|
|
app = create_search_assistant()
|
|
|
|
|
|
print("🔍 智能搜索助手启动！")
|
|
|
print("我会使用Tavily API为您搜索最新、最准确的信息")
|
|
|
print("支持各种问题：新闻、技术、知识问答等")
|
|
|
print("(输入 'quit' 退出)\n")
|
|
|
|
|
|
session_count = 0
|
|
|
|
|
|
while True:
|
|
|
user_input = input("🤔 您想了解什么: ").strip()
|
|
|
|
|
|
if user_input.lower() in ['quit', 'q', '退出', 'exit']:
|
|
|
print("感谢使用！再见！👋")
|
|
|
break
|
|
|
|
|
|
if not user_input:
|
|
|
continue
|
|
|
|
|
|
session_count += 1
|
|
|
config = {"configurable": {"thread_id": f"search-session-{session_count}"}}
|
|
|
|
|
|
# 初始状态
|
|
|
initial_state = {
|
|
|
"messages": [HumanMessage(content=user_input)],
|
|
|
"user_query": "",
|
|
|
"search_query": "",
|
|
|
"search_results": "",
|
|
|
"final_answer": "",
|
|
|
"step": "start"
|
|
|
}
|
|
|
|
|
|
try:
|
|
|
print("\n" + "="*60)
|
|
|
|
|
|
# 执行工作流
|
|
|
async for output in app.astream(initial_state, config=config):
|
|
|
for node_name, node_output in output.items():
|
|
|
if "messages" in node_output and node_output["messages"]:
|
|
|
latest_message = node_output["messages"][-1]
|
|
|
if isinstance(latest_message, AIMessage):
|
|
|
if node_name == "understand":
|
|
|
print(f"🧠 理解阶段: {latest_message.content}")
|
|
|
elif node_name == "search":
|
|
|
print(f"🔍 搜索阶段: {latest_message.content}")
|
|
|
elif node_name == "answer":
|
|
|
print(f"\n💡 最终回答:\n{latest_message.content}")
|
|
|
|
|
|
print("\n" + "="*60 + "\n")
|
|
|
|
|
|
except Exception as e:
|
|
|
print(f"❌ 发生错误: {e}")
|
|
|
print("请重新输入您的问题。\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
asyncio.run(main()) |