内存知识智能体 (Agent RAG)

import os
from dotenv import load_dotenv
from openai import OpenAI
from langchain_openai import ChatOpenAI

# 加载环境变量
load_dotenv()

# 配置 API Key (如未在 .env 设置，请在此处修改)
if not os.getenv("DASHSCOPE_API_KEY"):
    # os.environ["DASHSCOPE_API_KEY"] = "sk-xxxxxxxx"
    print("⚠️ 请设置 DASHSCOPE_API_KEY")

os.environ["DASHSCOPE_BASE_URL"] = "https://dashscope.aliyuncs.com/compatible-mode/v1"

# 初始化 LLM
llm = ChatOpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url=os.getenv("DASHSCOPE_BASE_URL"),
    model="qwen3-coder-plus",
    temperature=0,
)

# 初始化 OpenAI 客户端 (用于 Embedding)
client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url=os.getenv("DASHSCOPE_BASE_URL"),
)

from langchain_core.embeddings import Embeddings

class DashScopeEmbeddings(Embeddings):
    """DashScope 兼容的 Embeddings 封装。"""

    def __init__(self, model: str = "text-embedding-v4", dimensions: int = 1024):
        self.model = model
        self.dimensions = dimensions

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        vectors = []
        # 简单批处理，每次 10 条
        for i in range(0, len(texts), 10):
            chunk = texts[i : i + 10]
            try:
                response = client.embeddings.create(
                    model=self.model,
                    input=chunk,
                    dimensions=self.dimensions,
                )
                vectors.extend([item.embedding for item in response.data])
            except Exception as e:
                print(f"Embedding Error: {e}")
                vectors.extend([[0.0] * self.dimensions] * len(chunk))
        return vectors

    def embed_query(self, text: str) -> list[float]:
        response = client.embeddings.create(
            model=self.model,
            input=[text],
            dimensions=self.dimensions,
        )
        return response.data[0].embedding

import re
from pathlib import Path
from typing import Iterable
from langchain_core.documents import Document

def load_txt_documents(target_dir: Path) -> list[Document]:
    """读取目录下的 txt 文件并按空行分割为 Document。"""

    def split_on_blank(text: str) -> Iterable[str]:
        # 按dimensions维度分割（一版用于通识文档）
        # for block in re.split(rf"\n\s*{{{self.dimensions}}}\s*\n", text):
        
        # 按空行分割（兼容 Windows \r\n）
        for block in re.split(r"\n\s*\n", text):
            cleaned = block.strip()
            if cleaned:
                yield cleaned

    documents: list[Document] = []
    
    # 检查目录
    if not target_dir.exists():
        print(f"❌ 错误: 目录 {target_dir} 不存在！请确认当前路径。")
        return []

    files = sorted(target_dir.glob("*.txt"))
    if not files:
        print(f"❌ 警告: {target_dir} 下没有找到 .txt 文件")
        return []

    for path in files:
        print(f"📄 正在读取: {path.name}")
        content = path.read_text(encoding="utf-8")
        for idx, part in enumerate(split_on_blank(content)):
            documents.append(
                Document(
                    page_content=part,
                    metadata={"source": path.name, "chunk_id": idx},
                )
            )
    return documents

from langchain_core.vectorstores import InMemoryVectorStore

def build_vector_store() -> InMemoryVectorStore:
    # 指定 files 目录 (相对于当前 Notebook)
    target_dir = Path.cwd() / "files"
    documents = load_txt_documents(target_dir)

    embeddings = DashScopeEmbeddings()
    vector_store = InMemoryVectorStore(embedding=embeddings)
    
    if documents:
        vector_store.add_documents(documents)
        print(f"\n✅ 嵌入完成: 共 {len(documents)} 个片段。")
    
    return vector_store

# 立即构建
vector_store = build_vector_store()

from langchain.agents import create_agent
from langchain.tools import tool

def create_react_agent(vector_store: InMemoryVectorStore):
    """创建带检索工具的 ReAct Agent。"""

    @tool(response_format="content_and_artifact")
    def retrieve_context(query: str):
        """基于向量库检索与问题最相关的文本片段。"""
        # 打印日志方便调试
        print(f"   🔎 [检索中] Query: {query}")
        
        retrieved = vector_store.similarity_search(query, k=3)
        
        serialized = "\n\n".join(
            f"[{doc.metadata['source']}#{doc.metadata['chunk_id']}] {doc.page_content}"
            for doc in retrieved
        )
        return serialized, retrieved

    return create_agent(
        llm,
        tools=[retrieve_context],
        system_prompt=(
            "你可以使用检索工具获得参考资料。回答时结合检索到的内容，"
            "如有必要可以在答案中简单引用来源标识。"
        ),
    )

# 初始化 Agent
agent = create_react_agent(vector_store)

# === 修改此处的问题 ===
query = "考勤缺卡怎么处理？"

print(f"🚀 用户问题: {query}\n")

# 流式输出回答
for event in agent.stream({"messages": [{"role": "user", "content": query}]}, stream_mode="values"):
    event["messages"][-1].pretty_print()