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

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()

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

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

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

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")

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)
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()

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

print(f"關系數量: {len(relationship_df)}")
relationship_df.head()

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

print(f"報告記錄:{len(report_df)}")
report_df.head()

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

print(f"文本單元記錄:{len(text_unit_df)}")
text_unit_df.head()

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",  # 描述響應類型和格式的自由形式文本，可以是任何內容，例如優先列表、單段、多段、多頁報告
)

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

# 萊昂納多·達·芬奇
萊昂納多·達·芬奇，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）]。

總之，萊昂納多·達·芬奇對藝術、科學和工程的無與倫比的貢獻，加上他的創新思維和對同時代人及后代的深遠影響，使他成為人類成就史上的一位杰出人物。他的遺產繼續激發著欽佩和研究，強調了他的天才的永恒相關性。

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

["- 萊昂納多·達·芬奇的早期作品有哪些，它們存放在哪里？",
 "萊昂納多·達·芬奇與安德烈亞·德爾·維羅基奧的關系如何影響了他的早期作品？",
 "萊昂納多·達·芬奇在米蘭期間承擔了哪些重要項目？",
 "萊昂納多·達·芬奇的工程技能如何促成他的項目？",
 "萊昂納多·達·芬奇與法國弗朗西斯一世的關系有何重要性？"]

# import shutil
#
# shutil.rmtree(index_root)