llm = OpenAI(model="gpt-4o")
embed_model=OpenAIEmbedding(model_name="text-embedding-3-small"),

#加载数据
documents = SimpleDirectoryReader(input_files=['./dreamcity.txt']).load_data()

#图数据库
from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore
from llama_index.core import PropertyGraphIndex
from llama_index.core.indices.property_graph import SimpleLLMPathExtractor

#neo4j存储
graph_store = Neo4jPropertyGraphStore(
    username="neo4j",
    password="Unycp123!!",
    url="bolt://localhost:7687",
)

#知识抽取的提示词
prompt = '''
下面提供了一些文本。根据文本，提取最多 {max_knowledge_triplets} 个知识三元组，形式为（实体,关系,实体或属性）。避免使用停用词。仅输出三元组，不要有多余解释和说明。
---------------------
示例：
文本:Alice是Bob的母亲。
三元组:Alice,母亲,Bob
文本:Philz是1982年在伯克利创立的咖啡店。
三元组：
Philz,是,咖啡店
Philz,创立于,伯克利
Philz,创立于,1982年
--------------------
文本：{text}
三元组：
'''

#抽取结果的解析
def parse_fn(response_str: str) -> List[Tuple[str, str, str]]:
    lines = response_str.split("\n")
    triples = [line.split(",") for line in lines]
    return triples

#定义知识抽取器
kg_extractor = SimpleLLMPathExtractor(
    llm=llm,
    extract_prompt=prompt,
    max_paths_per_chunk=50,
    parse_fn=parse_fn,
)

#抽取创建图知识库索引（本地化做持久，避免反复抽取）
if not os.path.exists(f"./index_storage"):
    index = PropertyGraphIndex.from_documents(
        documents,
        embed_model=embed_model,
        kg_extractors=[kg_extractor],
        property_graph_store=graph_store,
        show_progress=True,
    )
    index.storage_context.persist("./index_storage")
else:

    print('Loading index...\n')
    storage_context = StorageContext.from_defaults(persist_dir="./index_storage",property_graph_store=graph_store)
    index = load_index_from_storage(storage_context=storage_context)

...
def parse_fn(output: str) -> list[str]:
    matches = output.strip().split("^")
    keywords=[x.strip().capitalize() for x in matches if x.strip()]
    print(keywords)
    return keywords

prompt = (
        "给定一些初始查询，提取最多 {max_keywords} 个相关关键词，考虑大小写、复数形式、常见表达等。\n"
        "用 '^' 符号分隔所有同义词/关键词：'关键词1^关键词2^...'\n"
        "注意，结果应为一行，用 '^' 符号分隔。"
        "----\n"
        "查询: {query_str}\n"
        "----\n"
        "关键词: "
    )

synonym_retriever = LLMSynonymRetriever(
    index.property_graph_store,
    llm=llm,
    include_text=False,
    output_parsing_fn=parse_fn,
    max_keywords=10,
    synonym_prompt=prompt,
    path_depth=1,
)

...
vector_retriever = VectorContextRetriever(
    index.property_graph_store,
    embed_model=embed_model,
    include_text=False,
    similarity_top_k=2,
    path_depth=1,
)

......
retriever = PGRetriever(sub_retrievers=[synonym_retriever,vector_retriever])
retrieve_nodes = retriever.retrieve("梅林发现的东西最后运用到哪里了？")
for node in retrieve_nodes:
    print(node.text)

...
query_engine = index.as_query_engine(
    sub_retrievers=[synonym_retriever,vector_retriever]
)

#查询
response = query_engine.query("梅林发现的东西最后运用到哪里了？")
print(response)