利用Milvus向量数据库，带你实现GraphRAG 原创

pip install --upgrade pymilvus
pip install git+https://github.com/zc277584121/graphrag.git
1.
2.

import nest_asyncio
nest_asyncio.apply()

import os
import urllib.request

index_root = os.path.join(os.getcwd(), 'graphrag_index')
os.makedirs(os.path.join(index_root, 'input'), exist_ok=True)
url = "https://www.gutenberg.org/cache/epub/7785/pg7785.txt"
file_path = os.path.join(index_root, 'input', 'davinci.txt')
urllib.request.urlretrieve(url, file_path)
with open(file_path, 'r+', encoding='utf-8') as file:
    # 使用文本文件的前934行，因为后面的行与本例无关。
    # 如果想节省API密钥成本，可以截断文本文件以减小大小。
    lines = file.readlines()
    file.seek(0)
    file.writelines(lines[:934])  # 如果想节省API密钥成本，可以减少这个数字。
    file.truncate()
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.

python -m graphrag.index --init --root ./graphrag_index
1.

python -m graphrag.index --root ./graphrag_index
1.

import os

import pandas as pd
import tiktoken
from graphrag.query.context_builder.entity_extraction import EntityVectorStoreKey
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.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.embedding import OpenAIEmbedding
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.question_gen.local_gen import LocalQuestionGen
from graphrag.query.structured_search.local_search.mixed_context import (
    LocalSearchMixedContext,
)
from graphrag.query.structured_search.local_search.search import LocalSearch
from graphrag.vector_stores import MilvusVectorStore
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.

output_dir = os.path.join(index_root, "output")
subdirs = [os.path.join(output_dir, d) for d in os.listdir(output_dir)]
latest_subdir = max(subdirs, key=os.path.getmtime)  # 获取最新的输出目录
INPUT_DIR = os.path.join(latest_subdir, "artifacts")
1.
2.
3.
4.

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
1.
2.
3.
4.
5.
6.
7.

# 读取节点表以获取社区和度量数据
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")
1.
2.
3.

entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)
description_embedding_store = MilvusVectorStore(
    collection_name="entity_description_embeddings",
)
# description_embedding_store.connect(uri="http://localhost:19530") # 用于Milvus docker服务
description_embedding_store.connect(uri="./milvus.db") # For Milvus Lite
entity_description_embeddings = store_entity_semantic_embeddings(
    entities=entities, vectorstore=description_embedding_store
)
print(f"实体数量:{len(entity_df)}")
entity_df.head()
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.

relationship_df = pd.read_parquet(f"{INPUT_DIR}/{RELATIONSHIP_TABLE}.parquet")
relationships = read_indexer_relationships(relationship_df)
1.
2.

print(f"关系数量: {len(relationship_df)}")
relationship_df.head()
1.
2.

report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
1.
2.

print(f"报告记录:{len(report_df)}")
report_df.head()
1.
2.

text_unit_df = pd.read_parquet(f"{INPUT_DIR}/{TEXT_UNIT_TABLE}.parquet")
text_units = read_indexer_text_units(text_unit_df)
1.
2.

print(f"文本单元记录:{len(text_unit_df)}")
text_unit_df.head()
1.
2.

api_key = os.environ["OPENAI_API_KEY"]  # 你的OpenAI API密钥
llm_model = "gpt-4o"  # 或 gpt-4-turbo-preview
embedding_model = "text-embedding-3-small"

llm = ChatOpenAI(
    api_key=api_key,
    model=llm_model,
    api_type=OpenaiApiType.OpenAI,
    max_retries=20,
)
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,
    relatinotallow=relationships,
    covariates=None, #covariates,#todo
    entity_text_embeddings=description_embedding_store,
    embedding_vectorstore_key=EntityVectorStoreKey.ID,  # 如果向量存储使用实体标题作为ID，则将此设置为EntityVectorStoreKey.TITLE
    text_embedder=text_embedder,
    token_encoder=token_encoder,
)

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,  # 根据你的模型的令牌限制更改此设置（如果你使用的是8k限制的模型，一个好设置可能是5000）
}

llm_params = {
    "max_tokens": 2_000,  # 根据你的模型的令牌限制更改此设置（如果你使用的是8k限制的模型，一个好设置可能是1000=1500）
    "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",  # 描述响应类型和格式的自由形式文本，可以是任何内容，例如优先列表、单段、多段、多页报告
)
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.

result = await search_engine.asearch("Tell me about Leonardo Da Vinci")
print(result.response)
1.
2.

# 莱昂纳多·达·芬奇
莱昂纳多·达·芬奇，1452年出生于佛罗伦萨附近的文奇镇，被广泛誉为意大利文艺复兴时期最多才多艺的天才之一。他的全名是莱昂纳多·迪·塞尔·皮耶罗·达·安东尼奥·迪·塞尔·皮耶罗·迪·塞尔·圭多·达·芬奇，他是塞尔·皮耶罗的非婚生和长子，塞尔·皮耶罗是一位乡村公证人[数据：实体（0]。莱昂纳多的贡献涵盖了艺术、科学、工程和哲学等多个领域，他被誉为基督教时代最万能的天才[数据：实体（8]。

## 早期生活和训练
莱昂纳多早期的才华得到了他父亲的认可，他父亲将他的一些画作带给了安德烈亚·德尔·维罗基奥，一位著名的艺术家和雕塑家。维罗基奥对莱昂纳多的才华印象深刻，于1469-1470年左右接受了他进入自己的工作室。在这里，莱昂纳多遇到了包括博蒂切利和洛伦佐·迪·克雷迪在内的其他著名艺术家[数据：来源（6, 7]。到1472年，莱昂纳多被佛罗伦萨画家行会录取，标志着他职业生涯的开始[数据：来源（7）]。

## 艺术杰作
莱昂纳多或许以他的标志性画作最为人所知，如《蒙娜丽莎》和《最后的晚餐》。《蒙娜丽莎》以其微妙的表情和详细的背景而闻名，现藏于卢浮宫，仍然是世界上最著名的艺术品之一[数据：关系（0, 45]。《最后的晚餐》是一幅壁画，描绘了耶稣宣布他的一个门徒将背叛他的那一刻，位于米兰圣玛利亚·格拉齐教堂的餐厅[数据：来源（2]。其他重要作品包括《岩间圣母》和他大约在1489-1490年开始的《绘画论》[数据：关系（7, 12]。

## 科学和工程贡献
莱昂纳多的天才超越了艺术，延伸到各种科学和工程事业。他在解剖学、光学和水力学方面做出了重要观察，他的笔记本里充满了预示许多现代发明的草图和想法。例如，他预示了哥白尼关于地球运动的理论和拉马克对动物的分类[数据：关系（38, 39]。他对光影法则和明暗对比的掌握对艺术和科学都产生了深远影响[数据：来源（45）]。

## 赞助和职业关系
莱昂纳多的职业生涯受到他的赞助人的重大影响。米兰公爵卢多维科·斯福尔扎雇佣莱昂纳多作为宫廷画家和普通工匠，委托了各种作品，甚至在1499年赠送给他一个葡萄园[数据：关系（9, 19, 84]。在他的晚年，莱昂纳多在法国国王弗朗西斯一世的赞助下搬到了法国，国王为他提供了丰厚的收入，并对他评价很高[数据：关系（114, 37]。莱昂纳多在法国安博瓦兹附近的克洛克斯庄园度过了他最后几年，国王经常拜访他，他得到了他的密友和助手弗朗切斯科·梅尔齐的支持[数据：关系（28, 122）]。

## 遗产和影响
莱昂纳多·达·芬奇的影响远远超出了他的一生。他在米兰创立了一所绘画学校，他的技术和教导被他的学生和追随者，如乔瓦尼·安布罗焦·达·普雷迪斯和弗朗切斯科·梅尔齐传承下去[数据：关系（6, 15, 28]。他的作品继续受到庆祝和研究，巩固了他作为文艺复兴时期最伟大的大师之一的遗产。莱昂纳多将艺术和科学融合的能力在两个领域都留下了不可磨灭的印记，激励着无数的艺术家和科学家[数据：实体（148, 86]; 关系（27, 12）]。

总之，莱昂纳多·达·芬奇对艺术、科学和工程的无与伦比的贡献，加上他的创新思维和对同时代人及后代的深远影响，使他成为人类成就史上的一位杰出人物。他的遗产继续激发着钦佩和研究，强调了他的天才的永恒相关性。
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.

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 Leonardo Da Vinci",
    "Leonardo's early works",
]

Generate questions based on history.

candidate_questions = await question_generator.agenerate(
        question_history=question_history, context_data=None, question_count=5
    )
candidate_questions.response

["- 莱昂纳多·达·芬奇的早期作品有哪些，它们存放在哪里？",
 "莱昂纳多·达·芬奇与安德烈亚·德尔·维罗基奥的关系如何影响了他的早期作品？",
 "莱昂纳多·达·芬奇在米兰期间承担了哪些重要项目？",
 "莱昂纳多·达·芬奇的工程技能如何促成他的项目？",
 "莱昂纳多·达·芬奇与法国弗朗西斯一世的关系有何重要性？"]
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.

# import shutil
#
# shutil.rmtree(index_root)
1.
2.
3.