RAG

from dotenv import load_dotenv
from langchain_community.embeddings import DashScopeEmbeddings
from sklearn.metrics.pairwise import cosine_similarity

# 加载环境变量
_ = load_dotenv()

# 用户问题
query = "过年要给不熟的亲戚发红包么？"

# 知识库
docs = [
    "不来往的人就不要给他钱",
    "海胆豆腐真好吃下次还吃",
    "半熟牛排淋上不熟的芝士",
]

# 初始化向量生成器
embeddings = DashScopeEmbeddings()

# 生成向量
qv = embeddings.embed_query(query)
dv = embeddings.embed_documents(docs)

# 计算余弦相似度
similarities = cosine_similarity([qv], dv)[0]
results = list(enumerate(similarities))
by_sim = sorted(results, key=lambda r: r[1], reverse=True)

# 余弦相似度大 -> 两个单位向量夹角小 -> 向量挨得更近
print("按余弦相似度排序：")
for i, s in by_sim:
    print('-', docs[i], s)

import os

# 配置 UA
MY_USER_AGENT = (
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 "
    "(KHTML, like Gecko) Version/17.0 Safari/605.1.15"
)
os.environ["USER_AGENT"] = MY_USER_AGENT

import bs4

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_community.document_loaders import WebBaseLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.vectorstores import InMemoryVectorStore
from langchain.tools import tool
from langchain.agents import create_agent

# 加载模型配置
_ = load_dotenv()

# 加载模型
llm = ChatOpenAI(
    model="qwen3-max",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url=os.getenv("DASHSCOPE_BASE_URL"),
)

# 加载文章内容
bs4_strainer = bs4.SoupStrainer(class_=("post"))
loader = WebBaseLoader(
    web_paths=("https://luochang212.github.io/posts/quick_bi_intro/",),
    bs_kwargs={"parse_only": bs4_strainer},
    requests_kwargs={"headers": {"User-Agent": MY_USER_AGENT}},
)
docs = loader.load()

assert len(docs) == 1

print(f"Total characters: {len(docs[0].page_content)}")
print(docs[0].page_content[:248])

# 文本分块
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,  # chunk size (characters)
    chunk_overlap=200,  # chunk overlap (characters)
    add_start_index=True,  # track index in original document
)
all_splits = text_splitter.split_documents(docs)

print(f"Split blog post into {len(all_splits)} sub-documents.")

# 初始化向量生成器
embeddings = DashScopeEmbeddings()

# 初始化内存向量存储
vector_store = InMemoryVectorStore(embedding=embeddings)

# 将文档添加到向量存储
document_ids = vector_store.add_documents(documents=all_splits)

print(document_ids[:2])

# 创建上下文检索工具
@tool(response_format="content_and_artifact")
def retrieve_context(query: str):
    """Retrieve information to help answer a query."""
    retrieved_docs = vector_store.similarity_search(query, k=2)
    serialized = "\n\n".join(
        (f"Source: {doc.metadata}\nContent: {doc.page_content}")
        for doc in retrieved_docs
    )
    return serialized, retrieved_docs

# 创建 ReAct Agent
agent = create_agent(
    llm,
    tools=[retrieve_context],
    system_prompt=(
        # If desired, specify custom instructions
        "You have access to a tool that retrieves context from a blog post. "
        "Use the tool to help answer user queries."
    )
)

# 调用 Agent
response = agent.invoke({
    "messages": [{"role": "user", "content": "当前的 Agent 能力有哪些局限性？"}]
})

# # 获取 Agent 的完整回复
# for message in result["messages"]:
#     message.pretty_print()

# 获取 Agent 的最终回复
response['messages'][-1].pretty_print()

# !pip install jieba

import jieba

from langchain_community.retrievers import BM25Retriever
from langchain_core.documents import Document

def chinese_tokenize(text: str) -> list[str]:
    """中文分词函数"""
    tokens = jieba.lcut(text)
    return [token for token in tokens if token.strip()]

# 1. 使用文本创建中文检索器
text_retriever = BM25Retriever.from_texts(
    [
        "何意味",
        "那很坏了",
        "这点小事也无所谓吧",
        "我替她原谅你了",
    ],
    k=2,
    preprocess_func=chinese_tokenize,
)

# 2. 使用文档创建中文检索器
doc_retriever = BM25Retriever.from_documents(
    [
        Document(page_content="辣椒炒肉拌面"),
        Document(page_content="肉蛋葱鸡"),
        Document(page_content="这下不熟了"),
        Document(page_content="铁串子"),
    ],
    k=2,
    preprocess_func=chinese_tokenize,
)

# 检索文本
text_retriever.invoke("一件小事")

# 检索文档
doc_retriever.invoke("拌面")

import random
from typing import List
from pydantic import BaseModel, Field

# 这是用户 query
query = "盘点海獭的黑历史"

# 这是 向量检索 召回的文本片段
dense_texts = [
    "某些海洋生物会乱扔垃圾",
    "海獭太可爱了",
    "海獭臭臭的",
]

# 这是 关键词检索 召回的文本片段
sparse_texts = [
    "海獭臭臭的",
    "雪鸮的黑历史",
]

# 定义 Agent 输出格式
class ReRankOutput(BaseModel):
    indices: List[int] = Field(description="重排后的召回文本片段的索引列表")

# 返回最多 limit 个文本片段
def get_relevant_texts(query: str,
                       dense_texts: list,
                       sparse_texts: list,
                       limit: int = 3):

    # 创建上下文
    texts = dense_texts + sparse_texts

    # 去重
    texts = list(set(texts))

    # 打乱元素顺序，消除由位置引入的 bias
    random.shuffle(texts)

    # 将索引 id 显式添加到文本片段前
    texts_with_index = [f"{i} - {text}" for i, text in enumerate(texts)]

    context = '\n\n'.join(texts_with_index)
    prompt = "\n".join([
        f"{context}",
        "---",
        "上面是RAG召回的多个文本片段。每个文本片段的格式为 [索引] - [内容]。",
        f"请返回最多{limit}个与用户问题有关的文本片段的索引（若相关性内容不足，允许少于{limit}条）。",
        "\n注意事项：",
        "1. 相关性更高的文本片段应该排在前面",
        "2. 返回的文本片段必须有助于回答用户问题！",
        f"\n用户问题：{query}",
        "文本片段的索引列表：",
    ])

    # 创建带结构化输出的Agent
    agent = create_agent(
        model=llm,
        system_prompt="你是一个召回文本相关性重排助手",
        response_format=ReRankOutput,
    )

    # 调用 Agent
    result = agent.invoke(
        {"messages": [{"role": "user", "content": prompt}]},
    )

    indices = result['structured_response'].indices
    return [texts[i] for i in indices]

res = get_relevant_texts(
    query,
    dense_texts,
    sparse_texts,
)

res