本地大模型编程实战(17)用langchain和智能体实现RAG(Retrieval Augmented Generation,检索增强生成)(1)

def create_retriever(embed_model_name):
    """创建检索器"""

    persist_directory = get_persist_directory(embed_model_name)
    db = LocalVectorDBChroma(embed_model_name,persist_directory)

    # 基于Chroma 的 vector store 生成 检索器
    vector_store = db.get_vector_store()
    retriever = vector_store.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 1},
    )
    return retriever

embed_model_name = "shaw/dmeta-embedding-zh"
retriever = create_retriever(embed_model_name)

def search(query):
    """查询矢量数据库"""

    persist_directory = get_persist_directory(embed_model_name)
    db = LocalVectorDBChroma(embed_model_name,persist_directory)
    vector_store = db.get_vector_store()

    results = vector_store.similarity_search_with_score(query,k=1)
    return results

system_prompt = """
您是问答任务的助手。
请使用以下**上下文**回答问题。如果您不知道答案，就说您不知道。
最多使用三句话并保持答案简洁。

下面请回答问题。

问题: {question}

上下文: {context}
"""

prompt = ChatPromptTemplate.from_messages([("system", system_prompt)])

def create_rag_chain(llm_model_name):
    """创建RAG链"""

    llm = ChatOllama(model=llm_model_name,temperature=0,verbose=True)
    rag_chain = (
        {
            "context": itemgetter("question") | retriever,
            "question": itemgetter("question")
        }
        | prompt
        | llm
        | StrOutputParser()
    )

    print(f'json_schema:{rag_chain.input_schema.model_json_schema()}')
    
    return rag_chain


def test_rag_chain(llm_model_name,question):
    """测试 rag 链"""

    print(f"------------{llm_model_name}-----------")
    rag_chain = create_rag_chain(llm_model_name)

    res = rag_chain.invoke({"question":question})
    print(res)

def test_rag_chain_stream(llm_model_name,question):
    """测试 rag 链，流式输出"""

    print(f"------------{llm_model_name}-----------")
    rag_chain = create_rag_chain(llm_model_name)
    for chunk in rag_chain.stream({"question":question}):
        print(chunk,end="|")

query = "猪的学名是什么？它对人类有什么用处？"

def create_agent(llm_model_name):
    """生成智能体"""

    rag_chain = create_rag_chain(llm_model_name)
    rag_tool = rag_chain.as_tool(
        name="animal_expert",
        description="获取有关动物的信息。",
    )

    llm = ChatOllama(model=llm_model_name,temperature=0,verbose=True)
    agent = create_react_agent(llm, [rag_tool])

    return agent

def test_agent(llm_model_name,question):

    agent = create_agent(llm_model_name)

    for chunk in agent.stream(
        {"messages": [("human",question)]}
    ):
        print(chunk)
        print("----")

query = "蜜蜂的特点是什么？它对人类社会有什么作用？"

{'agent': {'messages': [AIMessage(content='', ..., response_metadata={'model': 'qwen2.5', 'created_at': ...}, id=..., tool_calls=[{'name': 'animal_expert', 'args': {'question': '蜜蜂的特点是什么？它对人类社会有什么作用？'}, 'id': ..., 'type': 'tool_call'}], ...)]}}

{'tools': {'messages': [ToolMessage(content='蜜蜂的特点是勤劳且具有复杂的社会结构。它们对人类社会的作用主要体现在授粉、蜂蜜生产和提供蜂胶及蜂毒疗法等健康益处上。', name='animal_expert', id=..., tool_call_id=...)]}}

{'agent': {'messages': [AIMessage(content='蜜蜂是一种非常勤劳的昆虫，拥有复杂的社会结构。每只蜜蜂在蜂巢中都有其特定的角色和职责。\n\n蜜蜂对人类社会有着重要的作用：\n1. **授粉**：蜜蜂是自然界中最主要的传粉者之
一，对于维持生态平衡、保护生物多样性以及促进农作物产量具有重要作用。\n2. **蜂蜜生产**：通过采集花蜜酿造蜂蜜，不仅为人们提供了天然甜味剂和营养丰富的食品，还促进了相关产业的发展。\n3. **蜂胶及蜂毒疗法**：蜜蜂分 泌的蜂胶和蜂毒在传统医学中被用于治疗各种疾病，具有一定的药用价值。\n\n总之，蜜蜂对人类社会有着不可替代的作用。', additional_kwargs={}, response_metadata={'model': 'qwen2.5', ...}]}}