增强大模型高效检索：基于LlamaIndex ，构建一个轻量级带有记忆的 ColBERT 检索 Agent

在自然语言处理领域，高效检索相关信息的能力至关重要。将对话式记忆集成到文档检索系统中已经成为增强信息检索代理效果的强大技术。

在文中，我们专为 LlamaIndex 量身定制，将深入探讨构建一个轻量级的带有记忆的 ColBERT 检索代理，为高级检索任务提供简单而有效的解决方案。

我们还将探讨这种集成如何补充 ReAct 的功能，为 LlamaIndex 生态系统提供无缝的交互。

技术交流

节前，我们组织了一场算法岗技术&面试讨论会，邀请了一些互联网大厂朋友、今年参加社招和校招面试的同学。

针对大模型技术趋势、大模型落地项目经验分享、新手如何入门算法岗、该如何准备面试攻略、面试常考点等热门话题进行了深入的讨论。

总结链接如下：

• 《大模型面试宝典》(2024版) 正式发布！
• 《大模型实战宝典》(2024版) 正式发布！

前沿技术资讯、算法交流、求职内推、算法竞赛、面试交流(校招、社招、实习)等、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企开发者互动交流~

我们建了大模型算法岗技术与面试交流群， 想要进交流群、需要源码&资料、提升技术的同学，可以直接加微信号：mlc2060。加的时候备注一下：研究方向 +学校/公司+CSDN，即可。然后就可以拉你进群了。

方式①、微信搜索公众号：机器学习社区，后台回复：技术交流

方式②、添加微信号：mlc2060，备注：技术交流

定义

在我们踏上旅程之前，让我们澄清一些关键概念：

ColBERT：ColBERT，即基于 BERT 的上下文交互，是一种利用预训练语言模型如 BERT 来优化文档检索的技术。

HyDE：混合文档嵌入（HyDE）将稀疏和密集嵌入的优势结合起来，以获得更准确的文档表示。

对话式记忆：这指的是代理能够保留过去交互中的信息，从而实现更具上下文相关性的响应。

集成的好处

将对话式记忆集成到基于ColBERT的检索代理中带来了几个引人注目的优势：

• 上下文相关性：通过保留对话历史，代理可以根据正在进行的对话定制检索结果，提高相关性。
  连续性：对话式记忆促进了交互的连续性，使对话流程更加自然和连贯。
• 个性化：记忆集成使代理能够适应个人用户偏好和先前的交互，提升用户体验。

代码实现

让我们深入了解我们轻量级 ColBERT 检索代理与记忆的实现细节：

步骤 I：安装库

%pip install llama-index-core
%pip install llama-index-llms-openai
%pip install llama-index-embeddings-openai
%pip install llama-index-postprocessor-colbert-rerank
%pip install llama-index-readers-web

步骤 II：导入库，初始化 OpenAI，索引和加载数据

import os
from llama_index.readers.web import BeautifulSoupWebReader

from llama_index.core import VectorStoreIndex
from llama_index.embeddings.openai import OpenAIEmbedding 

from llama_index.core.query_pipeline import (
    QueryPipeline,
    InputComponent,
    ArgPackComponent,
)
from llama_index.core.prompts import PromptTemplate
from llama_index.llms.openai import OpenAI
from llama_index.postprocessor.colbert_rerank import ColbertRerank

from typing import Any, Dict, List, Optional
from llama_index.core.bridge.pydantic import Field
from llama_index.core.llms import ChatMessage
from llama_index.core.query_pipeline import CustomQueryComponent
from llama_index.core.schema import NodeWithScore

os.environ["OPENAI_API_KEY"] = "sk-..."

# 加载数据
reader = BeautifulSoupWebReader()

documents = reader.load_data(
    ["https://docs.anthropic.com/claude/docs/tool-use"]
)

# 索引
index = VectorStoreIndex.from_documents(
    documents,
    embed_model=OpenAIEmbedding(
        model="text-embedding-3-large", embed_batch_size=256
    ),
)

步骤 III：查询管道构建

# 首先，我们创建一个输入组件来捕获用户查询
input_component = InputComponent()

# 接下来，我们使用 LLM 重写用户查询
rewrite = (
    "请使用当前对话写一个查询给语义搜索引擎。\n"
    "\n"
    "\n"
    "{chat_history_str}"
    "\n"
    "\n"
    "最新消息：{query_str}\n"
    '查询："""\n'
)
rewrite_template = PromptTemplate(rewrite)
llm = OpenAI(
    model="gpt-4-turbo-preview",
    temperature=0.2,
)

# 我们将检索两次，因此需要将检索到的节点打包到一个列表中
argpack_component = ArgPackComponent()

# 使用它，我们将检索...
retriever = index.as_retriever(similarity_top_k=6)

# 然后使用 Colbert 进行后处理/重新排序
reranker = ColbertRerank(top_n=3)

步骤 VI：带有聊天历史的响应

DEFAULT_CONTEXT_PROMPT = (
    "以下是一些可能相关的上下文：\n"
    "-----\n"
    "{node_context}\n"
    "-----\n"
    "请使用上述上下文回答以下问题：\n"
    "{query_str}\n"
)


class ResponseWithChatHistory(CustomQueryComponent):
    llm: OpenAI = Field(..., description="OpenAI LLM")
    system_prompt: Optional[str] = Field(
        default=None, description="用于 LLM 的系统提示"
    )
    context_prompt: str = Field(
        default=DEFAULT_CONTEXT_PROMPT,
        description="用于 LLM 的上下文提示",
    )

    def _validate_component_inputs(
        self, input: Dict[str, Any]
    ) -> Dict[str, Any]:
        """在 run_component 期间验证组件输入。"""
        # 注意：这是可选的，但我们展示了在哪里进行验证作为示例
        return input

    @property
    def _input_keys(self) -> set:
        """输入键字典。"""
        # 注意：这些是必需的输入。如果有可选输入，请覆盖 `optional_input_keys_dict`
        return {"chat_history", "nodes", "query_str"}

    @property
    def _output_keys(self) -> set:
        return {"response"}

    def _prepare_context(
        self,
        chat_history: List[ChatMessage],
        nodes: List[NodeWithScore],
        query_str: str,
    ) -> List[ChatMessage]:
        node_context = ""
        for idx, node in enumerate(nodes):
            node_text = node.get_content(metadata_mode="llm")
            node_context += f"上下文块 {idx}：\n{node_text}\n\n"

        formatted_context = self.context_prompt.format(
            node_context=node_context, query_str=query_str
        )
        user_message = ChatMessage(role="user", content=formatted_context)

        chat_history.append(user_message)

        if self.system_prompt is not None:
            chat_history = [
                ChatMessage(role="system", content=self.system_prompt)
            ] + chat_history

        return chat_history

    def _run_component(self, **kwargs) -> Dict[str, Any]:
        """运行组件。"""
        chat_history = kwargs["chat_history"]
        nodes = kwargs["nodes"]
        query_str = kwargs["query_str"]

        prepared_context = self._prepare_context(
            chat_history, nodes, query_str
        )

        response = llm.chat(prepared_context)

        return {"response": response}

    async def _arun_component(self, **kwargs: Any) -> Dict[str, Any]:
        """异步运行组件。"""
        # 注意：可选的，但是异步 LLM 调用很容易实现
        chat_history = kwargs["chat_history"]
        nodes = kwargs["nodes"]
        query_str = kwargs["query_str"]

        prepared_context = self._prepare_context(
            chat_history, nodes, query_str
        )

        response = await llm.achat(prepared_context)

        return {"response": response}


response_component = ResponseWithChatHistory(
    llm=llm,
    system_prompt=(
        "你是一个问答系统。你将得到先前的聊天历史，"
        "以及可能相关的上下文，来协助回答用户消息。"
    ),
)
pipeline = QueryPipeline(
    modules={
        "input": input_component,
        "rewrite_template": rewrite_template,
        "llm": llm,
        "rewrite_retriever": retriever,
        "query

_retriever": retriever,
        "join": argpack_component,
        "reranker": reranker,
        "response_component": response_component,
    },
    verbose=False,
)

# 运行两次检索器 -- 一次使用虚构的查询，一次使用实际的查询
pipeline.add_link(
    "input", "rewrite_template", src_key="query_str", dest_key="query_str"
)
pipeline.add_link(
    "input",
    "rewrite_template",
    src_key="chat_history_str",
    dest_key="chat_history_str",
)
pipeline.add_link("rewrite_template", "llm")
pipeline.add_link("llm", "rewrite_retriever")
pipeline.add_link("input", "query_retriever", src_key="query_str")

# 每个输入到 argpack 组件都需要一个目标键 -- 它可以是任何东西
# 然后，argpack 组件将所有输入打包到一个列表中
pipeline.add_link("rewrite_retriever", "join", dest_key="rewrite_nodes")
pipeline.add_link("query_retriever", "join", dest_key="query_nodes")

# reranker 需要打包后的节点和查询字符串
pipeline.add_link("join", "reranker", dest_key="nodes")
pipeline.add_link(
    "input", "reranker", src_key="query_str", dest_key="query_str"
)

# synthesizer 需要重新排序后的节点和查询字符串
pipeline.add_link("reranker", "response_component", dest_key="nodes")
pipeline.add_link(
    "input", "response_component", src_key="query_str", dest_key="query_str"
)
pipeline.add_link(
    "input",
    "response_component",
    src_key="chat_history",
    dest_key="chat_history",
)

步骤 V：使用内存运行管道

from llama_index.core.memory import ChatMemoryBuffer

pipeline_memory = ChatMemoryBuffer.from_defaults(token_limit=8000)

user_inputs = [
    "你好！",
    "Claude-3 的工具使用是如何工作的？",
    "有哪些模型支持它？",
    "谢谢，这正是我需要知道的！",
]

for msg in user_inputs:
    # 获取记忆
    chat_history = pipeline_memory.get()

    # 准备输入
    chat_history_str = "\n".join([str(x) for x in chat_history])

    # 运行管道
    response = pipeline.run(
        query_str=msg,
        chat_history=chat_history,
        chat_history_str=chat_history_str,
    )

    # 更新记忆
    user_msg = ChatMessage(role="user", content=msg)
    pipeline_memory.put(user_msg)
    print(str(user_msg))

    pipeline_memory.put(response.message)
    print(str(response.message))
    print()

结论

将对话记忆集成到轻量级的 ColBERT 检索代理中，赋予其上下文意识和个性化交互能力。

通过按照为 LlamaIndex 定制的本文，你可以构建一个复杂而简化的检索系统，既提供相关信息，又保持对话连贯性。

拥抱增强记忆的检索能力，提升您在 LlamaIndex 生态系统中的自然语言处理应用。