import streamlit as st import os import json import re import datetime import tempfile # 导入 LangChain 相关组件 from langchain_openai import ChatOpenAI from langchain_core.prompts import ChatPromptTemplate from langchain_core.output_parsers import StrOutputParser from langchain_core.runnables import RunnableLambda from langchain_core.tools import tool from langchain_community.tools.tavily_search import TavilySearchResults from langchain_community.document_loaders import PyPDFLoader from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_huggingface import HuggingFaceEmbeddings from langchain_community.vectorstores import FAISS # 导入 YouTube 搜索库 from youtube_search import YoutubeSearch # ========================================== # 1. 基础配置与环境初始化 # ========================================== # 配置 Streamlit 页面标题、图标和布局 st.set_page_config(page_title="FeiChat Final", page_icon="✨", layout="wide") st.title("✨ FeiChat (Tavily + YouTube 完美版)") # 设置 API Key # 注意：实际生产中建议使用 st.secrets 或系统环境变量，不要直接写在代码里 os.environ["OPENAI_API_KEY"] = "lm-studio" # 指向本地 LM Studio，Key 随意填写 os.environ["TAVILY_API_KEY"] = "tvly-dev-xxxx" # Tavily 搜索引擎 Key # 初始化 Session State (会话状态) #用于在 Streamlit 页面刷新（rerun）时保存聊天记录和向量数据库 if "messages" not in st.session_state: st.session_state.messages = [] # 存储对话历史 if "vector_store" not in st.session_state: st.session_state.vector_store = None # 存储 PDF 向量索引 # ========================================== # 1.1 模型加载 (使用缓存避免重复加载) # ========================================== @st.cache_resource def get_models(): """ 初始化 LLM 和 Embedding 模型。 使用 @st.cache_resource 装饰器，确保只加载一次，节省资源。 """ # 1. 路由模型 (Router)：温度设为 0.0，要求输出精确，用于判断意图 router = ChatOpenAI( base_url="http://127.0.0.1:1234/v1", # 连接本地 LM Studio 端口 model="kuaidao-c-suite-v2", temperature=0.0 ) # 2. 对话模型 (Chat)：温度设为 0.7，用于生成流畅、自然的回答，开启流式输出 chat = ChatOpenAI( base_url="http://127.0.0.1:1234/v1", model="kuaidao-c-suite-v2", temperature=0.7, streaming=True ) # 3. 嵌入模型 (Embeddings)：用于将 PDF 文本转化为向量，这里使用 HuggingFace 的轻量级模型 embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2") return router, chat, embeddings llm_router, llm_chat, embeddings = get_models() # ========================================== # 2. 侧边栏 RAG (私有知识库处理) # ========================================== with st.sidebar: st.header("📂 私有知识库") # 文件上传控件 uploaded_file = st.file_uploader("上传 PDF (仅当问及文档内容时使用)", type=["pdf"]) # 如果用户上传了文件，且向量库还未建立，则开始处理 if uploaded_file and st.session_state.vector_store is None: with st.status("正在学习文档...", expanded=True): # 1. 创建临时文件保存上传的 PDF (PyPDFLoader 需要本地文件路径) with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: tmp.write(uploaded_file.read()) path = tmp.name # 2. 加载 PDF loader = PyPDFLoader(path) docs = loader.load() # 3. 文本切分：将长文档切成 500字符的小块，保留 50字符重叠以保持上下文 splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) splits = splitter.split_documents(docs) # 4. 向量化并存储：使用 FAISS 构建向量索引 st.session_state.vector_store = FAISS.from_documents(splits, embeddings) st.success(f"已索引 {len(splits)} 个片段") # 5. 删除临时文件，保持整洁 os.remove(path) # 清除记忆按钮 if st.button("🗑️ 清空记忆"): st.session_state.messages = [] st.session_state.vector_store = None st.rerun() # 重新运行脚本以刷新页面状态 # ========================================== # 3. 工具定义 (Search & RAG) # ========================================== # 初始化 Tavily 搜索客户端 tavily_engine = TavilySearchResults(max_results=5) @tool def internet_search(query: str) -> str: """Tavily 联网搜索工具函数，供 Agent 调用。""" print(f"🕵️ Tavily 搜索: {query}") # 后台打印日志 try: results = tavily_engine.invoke({"query": query}) formatted = [] # 格式化搜索结果，包含 URL 和内容摘要 for i, res in enumerate(results): formatted.append(f"【来源】({res['url']}):\n{res['content']}") return "\n\n".join(formatted) except Exception as e: return f"Error: {e}" @tool def knowledge_base_search(query: str) -> str: """知识库(RAG)搜索工具函数。""" # 如果没上传文件，直接返回提示 if st.session_state.vector_store is None: return "用户未上传任何文档。" # 在向量库中搜索最相似的 3 个片段 docs = st.session_state.vector_store.similarity_search(query, k=3) return "\n\n".join([f"【文档片段】: {d.page_content}" for d in docs]) # 将工具放入字典，方便 Router 调用 tools = {"internet_search": internet_search, "knowledge_base_search": knowledge_base_search} def search_youtube(query): """ YouTube 搜索辅助函数 注意：这是独立功能，不作为 LLM 的 Tool，而是在 UI 层直接展示结果 """ try: # 限制结果为 3 个 return YoutubeSearch(query, max_results=3).to_dict() except: return [] # ========================================== # 4. 核心逻辑 (Router & Chain) # ========================================== def get_time(): return datetime.datetime.now().strftime("%Y年%m月%d日") # 4.1 意图识别 Prompt # 作用：让 LLM 判断用户是想闲聊、查文档还是联网搜索，并提取搜索关键词 intent_prompt = ChatPromptTemplate.from_messages([ ("system", """ 你是一个智能路由。当前时间：{current_date}。 【工具选择逻辑】: 1. knowledge_base_search: 🔴 仅当用户明确提到“文档”、“PDF”、“上传的文件”等时使用。 2. internet_search: 🟢 默认选项（如果用户问知识性问题）。 3. CHAT: 仅用于纯打招呼、情感交流，不需要外部信息。 【Query生成规则】: - 强时效性问题（新闻、天气）：必须在关键词中加 `{current_date}`。 - 弱时效性问题（歌曲、百科、人物）：**禁止加日期**，直接用实体名。 返回 JSON 格式: {{ "intent": "CHAT" 或 "TOOL", "tool_name": "...", "tool_args": {{ "query": "..." }} }} """), ("user", "历史:\n{chat_history}\n\n输入:\n{input}") ]) def parse_router(text): """解析 Router LLM 返回的 JSON 字符串""" try: # 使用正则提取 Markdown 代码块中的 JSON (防止 LLM 输出 ```json ... ```) if "```" in text: text = re.search(r"```(?:json)?(.*?)```", text, re.DOTALL).group(1) return json.loads(text.strip()) except: # 解析失败则默认回退到纯聊天模式 return {"intent": "CHAT"} # 构建路由链：Prompt -> LLM -> 文本解析 -> JSON解析 intent_chain = intent_prompt | llm_router | StrOutputParser() | RunnableLambda(parse_router) # 4.2 最终回复润色 Prompt # 作用：根据搜索结果生成给用户的最终回答 response_prompt = ChatPromptTemplate.from_messages([ ("system", """ 你是一个严谨的信息整合助手。 请严格基于【搜索结果】回答。 1. 如果结果里没有，诚实说“未找到”。 2. 严禁捏造。 """), ("user", "问题: {user_input}\n\n搜索结果:\n{tool_result}") ]) response_chain = response_prompt | llm_chat | StrOutputParser() # 4.3 纯闲聊 Prompt chat_chain = ChatPromptTemplate.from_messages([ ("system", "助手。"), ("user", "{input}") ]) | llm_chat | StrOutputParser() # ========================================== # 5. 界面 UI 交互逻辑 # ========================================== # 5.1 显示历史消息 for msg in st.session_state.messages: with st.chat_message(msg["role"]): st.markdown(msg["content"]) # 5.2 处理用户输入 if user_input := st.chat_input("问：Fruits Zipper 最火的歌..."): # 记录用户输入 st.session_state.messages.append({"role": "user", "content": user_input}) with st.chat_message("user"): st.markdown(user_input) with st.chat_message("assistant"): # 使用 st.status 显示“思考中”状态动画 with st.status("🧠 思考中...", expanded=False) as status: # 准备上下文 hist = str(st.session_state.messages[:-1]) now = get_time() # --- 第一步：路由判断 --- intent_res = intent_chain.invoke({"input": user_input, "chat_history": hist, "current_date": now}) st.json(intent_res) # 调试用：在折叠状态里显示路由结果 final_stream = None # 用于存储最终的流式输出对象 yt_query = None # 用于存储 YouTube 搜索关键词 # --- 第二步：根据意图分支 --- if intent_res.get("intent") == "TOOL": tool_name = intent_res.get("tool_name") query = intent_res.get("tool_args", {}).get("query", user_input) # 如果是联网搜索，顺便记录一下关键词用于稍后搜 YouTube if tool_name == "internet_search": yt_query = query if tool_name in tools: try: # 执行工具 (Tavily 或 向量检索) tool_res = tools[tool_name].invoke(query) # 在 UI 中增加一个折叠框，显示搜索到的原始数据（增加可信度） with st.expander("📄 查看搜索摘要"): st.text(tool_res) # 防幻觉逻辑：如果搜索结果太短或包含错误信息 if "未找到有效信息" in tool_res or len(tool_res.strip()) < 50: final_stream = response_chain.stream({"user_input": user_input, "tool_result": "未找到相关信息。"}) else: # 正常生成回答 final_stream = response_chain.stream({"user_input": user_input, "tool_result": tool_res}) except Exception as e: st.error(f"工具执行失败: {e}") else: st.error("未找到工具") else: # 如果是 CHAT 意图，直接闲聊 final_stream = chat_chain.stream({"input": user_input}) # 更新状态栏为完成 status.update(label="完成", state="complete") # --- 第三步：流式输出回答 --- if final_stream: full_response = st.write_stream(final_stream) # Streamlit 自带的打字机效果 st.session_state.messages.append({"role": "assistant", "content": full_response}) # --- 第四步：展示 YouTube 视频 (仅当触发了联网搜索时) --- if yt_query: st.markdown("---") # 分割线 videos = search_youtube(yt_query) if videos: cols = st.columns(3) # 三列布局 for i, v in enumerate(videos[:3]): with cols[i]: # 处理缩略图：有些 API 返回的是列表，有些是字符串，这里做了兼容处理 thumb = v['thumbnails'][0] if isinstance(v['thumbnails'], list) else v['thumbnails'] st.image(thumb, use_container_width=True) # 显示标题链接 st.markdown(f"**[{v['title']}](https://www.youtube.com{v['url_suffix']})**") # 显示观看量 st.caption(f"👀 {v['views']}") else: st.caption("未找到相关视频。")