local search

引言

本地搜索（Local Search）是一种结合AI提取的知识图谱和原始文档文本块生成答案的方法。这种方法适用于需要理解文档中提到的特定实体的问题（例如，“洋甘菊的治愈特性是什么？”）。本文将详细介绍本地搜索的工作原理、实现细节以及如何在实际项目中应用它。

前置知识

在深入了解本地搜索之前，我们需要掌握以下几个关键概念：

1. 知识图谱：一种用于表示实体及其关系的图形数据结构。
2. LLM（Large Language Model）：大型语言模型，如GPT-4，能够理解和生成自然语言文本。
3. Pandas：一个强大的数据处理库，用于数据分析和操作。
4. Tiktoken：一个用于编码和解码文本的库，常用于处理语言模型的输入和输出。
5. LanceDB：一个用于存储和查询向量数据的库。

本地搜索的工作原理

本地搜索的核心思想是通过结合知识图谱中的结构化数据和输入文档中的非结构化数据，生成与特定实体相关的答案。具体步骤如下：

1. 加载数据：从parquet文件中加载实体、关系、协变量、社区报告和文本单元数据。
2. 构建上下文：使用本地搜索混合上下文构建器，将加载的数据转换为LLM提示的上下文。
3. 执行搜索：使用LLM生成答案，并返回结果。

代码示例

1. 设置LLM

首先，我们需要设置LLM模型。以下是代码示例：

import os
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.typing import OpenaiApiType

api_key = os.environ["GRAPHRAG_API_KEY"]
llm_model = os.environ["GRAPHRAG_LLM_MODEL"]

llm = ChatOpenAI(
    api_key=api_key,
    model=llm_model,
    api_type=OpenaiApiType.OpenAI,  # 可以是OpenaiApiType.OpenAI或OpenaiApiType.AzureOpenAI
    max_retries=20,
)

2. 加载数据

接下来，我们加载实体、关系、协变量、社区报告和文本单元数据。以下是代码示例：

import pandas as pd
from graphrag.query.indexer_adapters import (
    read_indexer_covariates,
    read_indexer_entities,
    read_indexer_relationships,
    read_indexer_reports,
    read_indexer_text_units,
)
from graphrag.query.input.loaders.dfs import store_entity_semantic_embeddings
from graphrag.vector_stores.lancedb import LanceDBVectorStore

INPUT_DIR = "./inputs/operation_dulce"
LANCEDB_URI = f"{INPUT_DIR}/lancedb"

COMMUNITY_REPORT_TABLE = "create_final_community_reports"
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"
RELATIONSHIP_TABLE = "create_final_relationships"
COVARIATE_TABLE = "create_final_covariates"
TEXT_UNIT_TABLE = "create_final_text_units"
COMMUNITY_LEVEL = 2

# 读取实体数据
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)

# 加载描述嵌入到内存中的LanceDB向量存储
description_embedding_store = LanceDBVectorStore(
    collection_name="entity_description_embeddings",
)
description_embedding_store.connect(db_uri=LANCEDB_URI)
entity_description_embeddings = store_entity_semantic_embeddings(
    entities=entities, vectorstore=description_embedding_store
)

# 读取关系数据
relationship_df = pd.read_parquet(f"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet")
relationships = read_indexer_relationships(relationship_df)

# 读取协变量数据
covariate_df = pd.read_parquet(f"{INPUT_DIR}/{COVARIATE_TABLE}.parquet")
claims = read_indexer_covariates(covariate_df)
covariates = {"claims": claims}

# 读取社区报告数据
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)

# 读取文本单元数据
text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet")
text_units = read_indexer_text_units(text_unit_df)

3. 构建本地搜索上下文

我们使用加载的数据构建本地搜索上下文。以下是代码示例：

from graphrag.query.structured_search.local_search.mixed_context import LocalSearchMixedContext
from graphrag.query.llm.oai.embedding import OpenAIEmbedding
import tiktoken

token_encoder = tiktoken.get_encoding("cl100k_base")

text_embedder = OpenAIEmbedding(
    api_key=api_key,
    api_base=None,
    api_type=OpenaiApiType.OpenAI,
    model=embedding_model,
    deployment_name=embedding_model,
    max_retries=20,
)

context_builder = LocalSearchMixedContext(
    community_reports=reports,
    text_units=text_units,
    entities=entities,
    relationships=relationships,
    covariates=covariates,  # 如果在索引期间没有运行协变量，设置为None
    entity_text_embeddings=description_embedding_store,
    embedding_vectorstore_key=EntityVectorStoreKey.ID,  # 如果向量存储使用实体标题作为ID，设置为EntityVectorStoreKey.TITLE
    text_embedder=text_embedder,
    token_encoder=token_encoder,
)

4. 创建本地搜索引擎

我们创建本地搜索引擎并执行搜索。以下是代码示例：

from graphrag.query.structured_search.local_search.search import LocalSearch

local_context_params = {
    "text_unit_prop": 0.5,
    "community_prop": 0.1,
    "conversation_history_max_turns": 5,
    "conversation_history_user_turns_only": True,
    "top_k_mapped_entities": 10,
    "top_k_relationships": 10,
    "include_entity_rank": True,
    "include_relationship_weight": True,
    "include_community_rank": False,
    "return_candidate_context": False,
    "embedding_vectorstore_key": EntityVectorStoreKey.ID,  # 如果向量存储使用实体标题作为ID，设置为EntityVectorStoreKey.TITLE
    "max_tokens": 12_000,  # 根据模型的token限制进行调整
}

llm_params = {
    "max_tokens": 2_000,  # 根据模型的token限制进行调整
    "temperature": 0.0,
}

search_engine = LocalSearch(
    llm=llm,
    context_builder=context_builder,
    token_encoder=token_encoder,
    llm_params=llm_params,
    context_builder_params=local_context_params,
    response_type="multiple paragraphs",  # 响应类型和格式，可以是任何形式
)

# 执行本地搜索
result = await search_engine.asearch("Tell me about Agent Mercer")
print(result.response)

5. 检查上下文数据

我们可以检查用于生成响应的上下文数据。以下是代码示例：

print(result.context_data["entities"].head())
print(result.context_data["relationships"].head())
print(result.context_data["reports"].head())
print(result.context_data["sources"].head())

if "claims" in result.context_data:
    print(result.context_data["claims"].head())

6. 问题生成

我们还可以使用本地搜索上下文生成候选问题。以下是代码示例：

from graphrag.query.question_gen.local_gen import LocalQuestionGen

question_generator = LocalQuestionGen(
    llm=llm,
    context_builder=context_builder,
    token_encoder=token_encoder,
    llm_params=llm_params,
    context_builder_params=local_context_params,
)

question_history = [
    "Tell me about Agent Mercer",
    "What happens in Dulce military base?",
]
candidate_questions = await question_generator.agenerate(
    question_history=question_history, context_data=None, question_count=5
)
print(candidate_questions.response)

结论

本地搜索是一种强大的工具，能够帮助我们从知识图谱和输入文档中生成与特定实体相关的答案。通过结合结构化数据和非结构化数据，本地搜索方法可以生成代表数据中最重要或最紧急信息内容或主题的答案。希望本文的详细解释和代码示例能够帮助你在实际项目中应用本地搜索技术。

参考资料

1. GraphRAG Documentation
2. OpenAI API Documentation
3. Pandas Documentation
4. Tiktoken Documentation
5. LanceDB Documentation

---

希望这篇文章能够帮助你更好地理解和使用本地搜索技术。如果你有任何问题或建议，欢迎在评论区留言！