传统RAG的实现

本文介绍基于 LangChain 的传统 RAG Agent 实现。

---

检索工具封装

将 Milvus 检索器封装为 Agent 可用的工具。

from langchain_core.tools import create_retriever_tool
from document.milvus_db import MilvusVectorSave

# 创建 Milvus 连接
mv = MilvusVectorSave()
mv.create_connection()

# 创建检索器
retriever = mv.vector_store_saved.as_retriever(
    search_type='similarity',
    search_kwargs={
        'k': 3,
        'score_threshold': 0.1,
        'ranker_type': 'rrf',
        'ranker_params': {'k': 100},
        'filter': {'category': 'content'}
    }
)

# 封装为 Agent 工具
retriever_tool = create_retriever_tool(
    retriever=retriever,
    name='rab_retriever',
    description='搜索并返回关于"半导体与芯片"的信息，内容涵盖：半导体与芯片的封装测试，光刻机,光刻胶等'
)

---

LLM 模型配置

配置支持 Agent 工具调用的 LLM 模型。

from langchain_openai import ChatOpenAI
from utils.env_utils import DEEPSEEK_API_KEY

llm = ChatOpenAI(
    temperature=0.5,
    model='deepseek-chat',
    api_key=DEEPSEEK_API_KEY,
    base_url='https://api.deepseek.com'
)

---

Agent 执行器构建

创建支持工具调用和记忆能力的 Agent。

from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory

# 创建提示词模板
prompt = ChatPromptTemplate.from_messages([
    ('system', '你是一个智能助手，你必须调用工具来回答用户的问题'),
    MessagesPlaceholder(variable_name='chat_history', optional=True),
    ('human', '{input}'),
    MessagesPlaceholder(variable_name='agent_scratchpad', optional=True),
])

# 创建 Agent
agent = create_tool_calling_agent(
    llm,
    [retriever_tool],
    prompt
)

# 创建执行器
executor = AgentExecutor(agent=agent, tools=[retriever_tool])

# 添加记忆能力
store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]

agent_with_history = RunnableWithMessageHistory(
    executor,
    get_session_history=get_session_history,
    input_messages_key='input',
    history_messages_key='chat_history'
)

---

使用示例

# 第一次对话
resp1 = agent_with_history.invoke(
    input={'input': '什么是EUV光刻机?'},
    config={'configurable': {'session_id': '111'}}
)
print(resp1)

# 第二次对话（带历史记忆）
resp2 = agent_with_history.invoke(
    input={'input': '你叫小莫?'},
    config={'configurable': {'session_id': '111'}}
)
print(resp2)

# 第三次对话
resp3 = agent_with_history.invoke(
    input={'input': '你是谁?'},
    config={'configurable': {'session_id': '111'}}
)
print(resp3)

---

完整代码汇总

# 导入必要的库
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.tools import create_retriever_tool
from langchain_openai import ChatOpenAI

import llm_models.all_llm_model
from document.milvus_db import MilvusVectorSave
from utils.env_utils import DEEPSEEK_API_KEY


# 1. 创建检索工具
def create_retriever_tool_func():
    mv = MilvusVectorSave()
    mv.create_connection()
    retriever = mv.vector_store_saved.as_retriever(
        search_type='similarity',
        search_kwargs={
            'k': 3,
            'score_threshold': 0.1,
            'ranker_type': 'rrf',
            'ranker_params': {'k': 100},
            'filter': {'category': 'content'}
        }
    )
    return create_retriever_tool(
        retriever=retriever,
        name='rab_retriever',
        description='搜索并返回关于"半导体与芯片"的信息'
    )


# 2. 创建 LLM
llm = ChatOpenAI(
    temperature=0.5,
    model='deepseek-chat',
    api_key=DEEPSEEK_API_KEY,
    base_url='https://api.deepseek.com'
)


# 3. 创建 Agent
retriever_tool = create_retriever_tool_func()

prompt = ChatPromptTemplate.from_messages([
    ('system', '你是一个智能助手，你必须调用工具来回答用户的问题'),
    MessagesPlaceholder(variable_name='chat_history', optional=True),
    ('human', '{input}'),
    MessagesPlaceholder(variable_name='agent_scratchpad', optional=True),
])

agent = create_tool_calling_agent(llm, [retriever_tool], prompt)
executor = AgentExecutor(agent=agent, tools=[retriever_tool])


# 4. 添加记忆
store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]

agent_with_history = RunnableWithMessageHistory(
    executor,
    get_session_history=get_session_history,
    input_messages_key='input',
    history_messages_key='chat_history'
)


# 5. 对话
if __name__ == '__main__':
    resp = agent_with_history.invoke(
        input={'input': '什么是EUV光刻机?'},
        config={'configurable': {'session_id': 'user_001'}}
    )
    print(resp)

---

流程图