LLM应用开发——AI搜索引擎实战

import os
from tavily import TavilyClient

class SearchEngine:
    """搜索引擎封装"""
    
    def __init__(self, api_key: str = None):
        self.client = TavilyClient(api_key=api_key or os.getenv("TAVILY_API_KEY"))
    
    def search(self, query: str, search_depth: str = "advanced", max_results: int = 5):
        """执行搜索"""
        # search_depth: "basic" 或 "advanced"
        response = self.client.search(
            query=query,
            search_depth=search_depth,
            max_results=max_results,
            include_answer=True, # Tavily 自带的简短回答
            include_raw_content=True # 包含抓取的网页原始内容
        )
        return response["results"]

# 使用示例
engine = SearchEngine()
results = engine.search("2024年诺贝尔物理学奖得主是谁？")

for res in results:
    print(f"Title: {res['title']}")
    print(f"URL: {res['url']}")
    print(f"Snippet: {res['content'][:100]}...")

import requests

def get_markdown_content(url: str):
    """使用 Jina Reader 提取网页内容"""
    reader_url = f"https://r.jina.ai/{url}"
    response = requests.get(reader_url)
    return response.text

prompt = """
你是一个专业的搜索助手。请根据提供的参考资料回答用户问题。

参考资料：
{context}

要求：
1. 必须在回答中引用参考资料，格式为 [n]。
2. 如果资料中没有相关信息，请直说不知道。
3. 回答要客观、准确。

用户问题：{query}
"""

# context 示例：
# [1] 来源: https://example.com 内容: 2024年诺贝尔物理学奖授予了...
# [2] 来源: https://news.com 内容: 瑞典皇家科学院宣布...

import requests

def extract_content(url: str) -> str:
    """使用 Jina Reader 提取网页内容"""
    reader_url = f"https://r.jina.ai/{url}"
    response = requests.get(reader_url)
    if response.status_code == 200:
        return response.text
    return ""

# 示例
content = extract_content("https://en.wikipedia.org/wiki/Large_language_model")
print(content[:500])

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser

class QueryOptimizer:
    """查询优化器"""
    
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o-mini")
        self.parser = JsonOutputParser()
    
    def optimize(self, query: str) -> list[str]:
        """将用户问题分解为多个搜索词"""
        prompt = ChatPromptTemplate.from_template("""
        你是一个搜索专家。请将用户的复杂问题分解为 2-3 个独立的、适合搜索引擎的关键词或短语。
        
        用户问题：{query}
        
        请以 JSON 数组格式输出，例如：["关键词1", "关键词2"]
        """)
        
        chain = prompt | self.llm | self.parser
        return chain.invoke({"query": query})

# 示例
optimizer = QueryOptimizer()
search_queries = optimizer.optimize("对比 Llama 3.1 和 GPT-4o 在代码生成方面的差异")
# 输出: ["Llama 3.1 vs GPT-4o code generation benchmarks", "Llama 3.1 coding capabilities", "GPT-4o programming performance"]

def generate_answer_with_citations(query: str, contexts: list[dict]):
    """生成带引用的回答"""
    
    # 构建上下文字符串
    context_str = ""
    for i, res in enumerate(contexts):
        context_str += f"Source [{i+1}]:\nTitle: {res['title']}\nURL: {res['url']}\nContent: {res['content']}\n\n"
    
    prompt = f"""
    你是一个专业的 AI 搜索助手。请根据提供的参考资料回答用户的问题。
    
    要求：
    1. 必须在回答中通过 [n] 的形式标注引用来源，其中 n 是参考资料的编号。
    2. 如果资料中没有相关信息，请诚实说明。
    3. 回答要客观、准确、结构清晰。
    4. 在回答末尾列出所有参考资料的链接。
    
    参考资料：
    {context_str}
    
    用户问题：{query}
    
    回答：
    """
    
    llm = ChatOpenAI(model="gpt-4o")
    return llm.invoke(prompt).content

from typing import Annotated, List, TypedDict
from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from tavily import TavilyClient

# 1. 定义状态
class AgentState(TypedDict):
    query: str
    search_queries: List[str]
    search_results: List[dict]
    answer: str
    steps: int

# 2. 定义节点逻辑
class SearchAgent:
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o")
        self.tavily = TavilyClient()

    async def plan_search(self, state: AgentState):
        """根据问题生成搜索计划"""
        prompt = f"针对问题 '{state['query']}'，生成 3 个互补的搜索关键词。"
        res = await self.llm.ainvoke(prompt)
        return {"search_queries": res.content.split("\n"), "steps": state.get("steps", 0) + 1}

    async def execute_search(self, state: AgentState):
        """执行并行搜索"""
        results = []
        for q in state["search_queries"]:
            res = self.tavily.search(q, max_results=3)
            results.extend(res["results"])
        return {"search_results": results}

    async def synthesize(self, state: AgentState):
        """综合信息生成回答"""
        context = "\n".join([r["content"] for r in state["search_results"]])
        prompt = f"基于以下内容回答 '{state['query']}'：\n\n{context}"
        res = await self.llm.ainvoke(prompt)
        return {"answer": res.content}

# 3. 构建图
workflow = StateGraph(AgentState)
agent = SearchAgent()

workflow.add_node("plan", agent.plan_search)
workflow.add_node("search", agent.execute_search)
workflow.add_node("synthesize", agent.synthesize)

workflow.set_entry_point("plan")
workflow.add_edge("plan", "search")
workflow.add_edge("search", "synthesize")
workflow.add_edge("synthesize", END)

app = workflow.compile()

def hyde_search(query: str):
    # 1. 生成假设性回答
    hypothetical_answer = llm.invoke(f"请针对问题 '{query}' 写一个简短的假设性回答。")
    
    # 2. 使用假设性回答进行向量检索或网页搜索
    # 这样可以显著提升检索的召回率 (Recall)
    results = tavily.search(hypothetical_answer.content)
    return results

def verify_citations(answer: str, sources: List[dict]):
    """
    校验生成的回答中的每一条引用是否在来源中真实存在。
    """
    prompt = f"请检查以下回答中的引用标注是否准确。如果来源中没有提到相关事实，请指出。\n\n回答：{answer}\n\n来源：{sources}"
    verification = llm.invoke(prompt)
    return verification.content

from pydantic import BaseModel, Field

class Citation(BaseModel):
    source_id: int
    quote: str = Field(description="引用的原文片段")

class SearchResponse(BaseModel):
    answer: str = Field(description="生成的回答内容")
    citations: List[Citation]
    references: List[str] = Field(description="参考链接列表")