MATCH (c:__Community__)
RETURN c.level AS level, COUNT(*) AS count
ORDER BY level

python -m graphrag.query 
--root ./msgraphrag 
--method local 
--community_level 2 
'X教授的名字是怎么来的？他和天启之间有怎样的联系'

python -m graphrag.query
--root ./msgraphrag
--method global
--community_level 2
'请概括介绍文章中复仇者联盟的相关信息'

CREATE VECTOR INDEX entity_index IF NOT EXISTS FOR (e:__Entity__) ON e.description_embedding
OPTIONS {indexConfig: {
 `vector.dimensions`: 1536,
 `vector.similarity_function`: 'cosine'
}}

......
os.environ["AZURE_OPENAI_API_KEY"] = '自行准备'
os.environ["AZURE_OPENAI_ENDPOINT"] = '自行准备'

#模型准备，后续也会用到
text_embedder = AzureOpenAIEmbeddings(
    azure_deployment='text-embedding-3-small',
    openai_api_version='2024-05-01-preview',
)

llm = AzureChatOpenAI(
    deployment_name='gpt-4o-mini',
    openai_api_version='2024-05-01-preview',
)

#以下为测试代码，实际不需要：使用Langchain的Neo4jVector组件，从已经创建的neo4j中的向量索引进行检索测试
entity_vector = Neo4jVector.from_existing_index(
    text_embedder,
    url=NEO4J_URI,  
    username=NEO4J_USERNAME,
    password=NEO4J_PASSWORD,
    index_name='entity_index',
    text_node_property='description',  #这个不可少，用来指定节点的text属性字段
)

result = entity_vector.similarity_search("复仇者联盟",top_k=5)
print(result[0].page_content)

lc_retrieval_query = """

    //接收向量检索输出的node，在此基础上进一步检索
    WITH collect(node) as nodes

    //查找最相关的文本块，输出text属性
    WITH
    collect {
        UNWIND nodes as n
        MATCH (n)<-[:HAS_ENTITY]->(c:__Chunk__)
        WITH c, count(distinct n) as freq
        RETURN c.text AS chunkText
        ORDER BY freq DESC
        LIMIT $topChunks
    } AS text_mapping,

    //查找最相关的社区，输出summary摘要（如果没有weight，用cypher设定）
    collect {
        UNWIND nodes as n
        MATCH (n)-[:IN_COMMUNITY]->(c:__Community__)
        WITH c, c.rank as rank, c.weight AS weight
        RETURN c.summary 
        ORDER BY rank, weight DESC
        LIMIT $topCommunities
    } AS report_mapping,

    //查找最相关的其他实体（nodes外部），输出描述
    collect {
        UNWIND nodes as n
        MATCH (n)-[r:RELATED]-(m) 
        WHERE NOT m IN nodes
        RETURN r.description AS descriptionText
        ORDER BY r.rank, r.weight DESC 
        LIMIT $topOutsideRels
    } as outsideRels,
      
    //查找最相关的其他实体（nodes内部），输出描述
    collect {
         UNWIND nodes as n
         MATCH (n)-[r:RELATED]-(m) 
         WHERE m IN nodes
         RETURN r.description AS descriptionText
         ORDER BY r.rank, r.weight DESC 
         LIMIT $topInsideRels
     } as insideRels,

     //输出这些实体本身的描述
     collect {
         UNWIND nodes as n
         RETURN n.description AS descriptionText
     } as entities

//返回text，score，metadata三个字段
RETURN {Chunks: text_mapping, Reports: report_mapping, 
       Relationships: outsideRels + insideRels, 
       Entities: entities} AS text, 1.0 AS score, {source:''} AS metadata 
"""

#创建neo4j向量存储对象，注意传入retrieval_query参数
lc_vector = Neo4jVector.from_existing_index(
    text_embedder,
    url=NEO4J_URI,
    username=NEO4J_USERNAME,
    password=NEO4J_PASSWORD,
    index_name='entity_index',
    retrieval_query=lc_retrieval_query,
)

#chain,并调用获得响应。此处可参考Langchain文档学习
chain = RetrievalQAWithSourcesChain.from_chain_type(
    llm,
    chain_type="stuff",         retriever=lc_vector.as_retriever(search_kwargs={"params":{            "topChunks": topChunks,
            "topCommunities": topCommunities,
            "topOutsideRels": topOutsideRels,
            "topInsideRels": topInsideRels,
        }})
)

response = chain.invoke(
    {"question": "复仇者联盟与钢铁侠有什么关系？"},
    return_only_outputs=True,
)
print(response['answer'])