利用PebbloRetrievalQA实现身份感知的RAG-CSDN博客

在当今的数据驱动时代，提供安全、精确的信息检索至关重要。PebbloRetrievalQA代表了这种提升安全性和准确性的尝试，它是一个结合身份和语义约束的检索链，用于从向量数据库中进行问答。本次分享将深入探讨如何利用PebbloRetrievalQA实现对文档的身份和语义强制检索，并提供可直接运行的代码示例。

技术背景介绍

PebbloRetrievalQA是一种检索链，通过在问答过程中对身份和语义进行约束，确保信息的请求和返回是安全且符合授权的。这一功能特别适用于需要高等级数据保密和合规性的应用场景，如金融报告、员工记录等。

核心原理解析

Pebblo提供身份和语义约束的功能：

身份约束: 通过身份验证确保只有被授权的用户能够访问特定信息。
语义约束: 限制那些与特定主题或实体无关的内容，从而保持信息的相关性和保密性。

代码实现演示

我们将使用PebbloRetrievalQA加载文档并实现身份和语义执行。

1. 加载文档并进行身份验证

首先，我们需要将文档加载到内存中的Qdrant向量数据库中，该数据库将作为PebbloRetrievalQA的检索工具。

from langchain_community.vectorstores.qdrant import Qdrant
from langchain_core.documents import Document
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_openai.llms import OpenAI

# 初始化语言模型和嵌入
llm = OpenAI()
embeddings = OpenAIEmbeddings()
collection_name = "pebblo-identity-and-semantic-rag"

page_content = """
... (省略的报告内容)
"""

documents = [
    Document(
        **{
            "page_content": page_content,
            "metadata": {
                "pebblo_semantic_topics": ["financial-report"],
                "pebblo_semantic_entities": ["us-bank-account-number"],
                "authorized_identities": ["finance-team", "exec-leadership"],
                "page": 0,
                "source": "https://drive.google.com/file/d/xxxxxxxxxxxxx/view",
                "title": "ACME Corp Financial Report.pdf",
            },
        }
    )
]

# 加载文档到Qdrant
vectordb = Qdrant.from_documents(
    documents,
    embeddings,
    location=":memory:",
    collection_name=collection_name,
)

print("Vectordb loaded.")

2. 实现身份检索

接下来，我们使用PebbloRetrievalQA来实现身份约束。

from langchain_community.chains import PebbloRetrievalQA
from langchain_community.chains.pebblo_retrieval.models import AuthContext, ChainInput

# 初始化PebbloRetrievalQA链
qa_chain = PebbloRetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectordb.as_retriever(),
    app_name="pebblo-identity-rag",
    description="Identity Enforcement app using PebbloRetrievalQA",
    owner="ACME Corp",
)

def ask(question: str, auth_context: dict):
    """
    询问PebbloRetrievalQA链问题
    """
    auth_context_obj = AuthContext(**auth_context) if auth_context else None
    chain_input_obj = ChainInput(query=question, auth_context=auth_context_obj)
    return qa_chain.invoke(chain_input_obj.dict())

# 授权用户提问
auth = {
    "user_id": "finance-user@acme.org",
    "user_auth": [
        "finance-team",
    ],
}
question = "Share the financial performance of ACME Corp for the year 2020"
resp = ask(question, auth)
print(f"Question: {question}\n\nAnswer: {resp['result']}")

3. 实现语义检索

除了身份验证外，我们还可以使用语义约束机制。

from typing import List, Optional
from langchain_community.chains import PebbloRetrievalQA
from langchain_community.chains.pebblo_retrieval.models import ChainInput, SemanticContext

# 初始化语义PebbloRetrievalQA链
qa_chain = PebbloRetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectordb.as_retriever(),
    app_name="pebblo-semantic-rag",
    description="Semantic Enforcement app using PebbloRetrievalQA",
    owner="ACME Corp",
)

def ask(question: str, topics_to_deny: Optional[List[str]] = None, entities_to_deny: Optional[List[str]] = None):
    """
    询问PebbloRetrievalQA链问题
    """
    semantic_context = dict()
    if topics_to_deny:
        semantic_context["pebblo_semantic_topics"] = {"deny": topics_to_deny}
    if entities_to_deny:
        semantic_context["pebblo_semantic_entities"] = {"deny": entities_to_deny}

    semantic_context_obj = SemanticContext(**semantic_context) if semantic_context else None
    chain_input_obj = ChainInput(query=question, semantic_context=semantic_context_obj)
    return qa_chain.invoke(chain_input_obj.dict())

# 测试语义约束
topic_to_deny = ["financial-report"]
entities_to_deny = []
question = "Share the financial performance of ACME Corp for the year 2020"
resp = ask(question, topics_to_deny=topic_to_deny, entities_to_deny=entities_to_deny)
print(f"Topics to deny: {topic_to_deny}\nEntities to deny: {entities_to_deny}\nQuestion: {question}\nAnswer: {resp['result']}")