Week 7-8：RAG（检索增强生成）

# 例子：用 Embedding 模型把文字转成向量
"BTC 价格突破 10 万" → [0.12, -0.34, 0.56, ..., 0.78]  # 768 维
"比特币创下历史新高" → [0.11, -0.32, 0.55, ..., 0.76]  # 向量很接近 ✨
"今天吃了一碗牛肉面" → [0.01, 0.02, -0.10, ..., 0.05]  # 向量完全不同

# app/rag/rag_pipeline.py
"""
RAG Pipeline 示例：
把 electric-trading-ai 项目的文档喂进向量库，然后问问题
"""

from pathlib import Path
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader, DirectoryLoader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

# ===== 配置 =====
PERSIST_DIR = "./chroma_db"              # 向量数据库持久化目录
DOCS_DIR = Path("e:/wf_work/wf_ai/electric-trading-ai")  # 你的项目根目录

# ===== 1. 加载和分块文档 =====

def load_and_split_docs(docs_dir: Path):
    """加载项目里的文档，然后分块"""
    # 加载 Markdown、txt、py 文件（选择性排除）
    loader = DirectoryLoader(
        str(docs_dir),
        glob=["**/*.md", "**/*.txt"],
        recursive=True,
        use_multithreading=True,
    )
    docs = loader.load()
    print(f"加载了 {len(docs)} 个文档")

    # 分块：长文档切成小段落
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,          # 每块大约 500 字符
        chunk_overlap=50,        # 相邻块有 50 字符重叠，避免切断语义
        length_function=len,
    )
    splits = text_splitter.split_documents(docs)
    print(f"切分成了 {len(splits)} 个块")
    return splits


# ===== 2. Embedding + 存进 ChromaDB =====

def build_vector_store(splits):
    """用 Embedding 模型向量化并存入 ChromaDB"""
    # 本地 Embedding 模型（支持中文，首次会自动下载约 400MB）
    embeddings = HuggingFaceEmbeddings(
        model_name="shibing624/text2vec-base-chinese",  # 中文效果好
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True},
    )

    # 创建/加载 ChromaDB
    vectorstore = Chroma.from_documents(
        documents=splits,
        embedding=embeddings,
        persist_directory=PERSIST_DIR,
        collection_name="trading_docs",
    )
    print(f"向量数据库已就绪，持久化到 {PERSIST_DIR}")
    return vectorstore


# ===== 3. 构建 RAG Chain =====

def build_rag_chain(vectorstore):
    """构建 RAG 问答链"""
    # 检索器：根据问题找 Top-4 最相关的文档块
    retriever = vectorstore.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 4},
    )

    # LLM
    llm = ChatOpenAI(
        model="Qwen/Qwen2.5-7B-Instruct",
        base_url="https://api.siliconflow.cn/v1",
        api_key="your-api-key-here",
        temperature=0.3,  # RAG 场景要低温度，减少发散
    )

    # Prompt 模板
    prompt = ChatPromptTemplate.from_template("""
你是一个 AI Agent 学习助手。请严格基于以下上下文回答用户的问题。
如果上下文中没有相关信息，请说"抱歉，我在你的文档里没有找到相关内容"，不要瞎编。

上下文：
{context}

用户问题：{question}
""")

    # LangChain LCEL Chain 组装
    rag_chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )

    return rag_chain


def format_docs(docs) -> str:
    """把检索到的文档块拼成一个字符串"""
    return "\n\n".join(doc.page_content for doc in docs)


# ===== 4. 运行 =====

async def main():
    # Step 1-2: 构建向量库（第一次跑会慢，之后向量库持久化了就快）
    splits = load_and_split_docs(DOCS_DIR)
    vectorstore = build_vector_store(splits)

    # Step 3: 构建 RAG Chain
    rag_chain = build_rag_chain(vectorstore)

    # Step 4: 问答
    questions = [
        "这个 electric-trading-ai 项目有哪些依赖？",
        "项目用的什么框架？",
        "Python 版本要求是多少？",
        "你好",  # 测试无关问题
    ]

    for q in questions:
        print(f"\n{'='*50}")
        print(f"Q: {q}")
        answer = await rag_chain.ainvoke(q)
        print(f"A: {answer}")


if __name__ == "__main__":
    import asyncio
    asyncio.run(main())

from langchain_core.tools import tool

# 把 RAG 检索器包装成一个 Tool
@tool
def query_project_docs(question: str) -> str:
    """当用户问到关于本项目的技术问题时，用这个工具去查项目文档。"""
    # 这里的 vectorstore 就是上一节构建的
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    docs = retriever.invoke(question)
    context = "\n\n".join(d.page_content for d in docs)
    return f"检索到以下相关内容：\n{context}"

# 然后把 query_project_docs 加进 Month 2 Week 5 的 TOOLS 列表
# Agent 就会自动在需要的时候调用它！