require python==3.12
pip install graphrag==0.5.0

!mkdir -p ./ragtest4/input/

graphrag init --root ./ragtest4

python -m graphrag prompt-tune --config ./settings.yaml --root ./ --language Chinese --output ./prompts2

graphrag index --root ./ragtest4

documents_df = pd.read_parquet("./output/create_final_documents.parquet")
documents_df

graphrag query --root ./ragtest4 --method local --query "请介绍下贾母"   #local模式

!graphrag query --root ./ragtest4 --method global --query "请介绍下贾母" # global模式

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_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.lancedb import LanceDBVectorStore
from IPython.display import Markdown, display

INPUT_DIR = "./ragtest4/output"
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"
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 = LanceDBVectorStore(
    collection_name="default-entity-description",
)
description_embedding_store.connect(db_uri=LANCEDB_URI)
entity_description_embeddings = store_entity_semantic_embeddings(
    entities=entities, vectorstore=description_embedding_store
)

print(f"Entity count: {len(entity_df)}")


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

print(f"Relationship count: {len(relationship_df)}")


report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
print(reports[0])
print(f"Report records: {len(report_df)}")



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

print(f"Text unit records: {len(text_unit_df)}")




api_key = "**********************"
llm_model = "glm-4-flash"
embedding_model = "nomic-embed-text"

llm = ChatOpenAI(
    api_key=api_key,
    model=llm_model,
    api_base="https://open.bigmodel.cn/api/paas/v4/",
    api_type=OpenaiApiType.OpenAI,  
    max_retries=20,
)

token_encoder = tiktoken.get_encoding("cl100k_base")

text_embedder = OpenAIEmbedding(
    api_key=api_key,
    api_base="http://localhost:11434/api",
    api_type=OpenaiApiType.OpenAI,
    model=embedding_model,
    max_retries=20,
)

context_builder = LocalSearchMixedContext(
    community_reports=reports,
    text_units=text_units,
    entities=entities,
    relationships=relationships,
    covariates=None,
    entity_text_embeddings=description_embedding_store,
    embedding_vectorstore_key=EntityVectorStoreKey.ID,  
    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": True,
    "return_candidate_context": True,
    "embedding_vectorstore_key": EntityVectorStoreKey.ID,  
    "max_tokens": 12_000, 
}

llm_params = {
    "max_tokens": 2_000, 
    "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("贾府有多少人")
print(result.response)
# 很抱歉，您提供的数据表中没有包含关于贾府人数的具体信息。
# 因此，我无法根据现有数据回答贾府有多少人。
# 如果您能提供相关的数据或背景信息，我将很乐意帮助您分析。

from graphrag.query.indexer_adapters import (
    read_indexer_communities,
    read_indexer_entities,
    read_indexer_reports,
)
from graphrag.query.llm.oai.chat_openai import ChatOpenAI
from graphrag.query.llm.oai.typing import OpenaiApiType
from graphrag.query.structured_search.global_search.community_context import (
    GlobalCommunityContext,
)
from graphrag.query.structured_search.global_search.search import GlobalSearch

COMMUNITY_TABLE = "create_final_communities"
COMMUNITY_REPORT_TABLE = "create_final_community_reports"
ENTITY_TABLE = "create_final_nodes"
ENTITY_EMBEDDING_TABLE = "create_final_entities"


community_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_TABLE}.parquet")
entity_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_TABLE}.parquet")
report_df = pd.read_parquet(f"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet")
entity_embedding_df = pd.read_parquet(f"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet")

communities = read_indexer_communities(community_df, entity_df, report_df)
reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)
entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)

context_builder = GlobalCommunityContext(
    community_reports=reports,
    communities=communities,
    entities=entities,  
    token_encoder=token_encoder,
)


context_builder_params = {
    "use_community_summary": False,  
    "shuffle_data": True,
    "include_community_rank": True,
    "min_community_rank": 0,
    "community_rank_name": "rank",
    "include_community_weight": True,
    "community_weight_name": "occurrence weight",
    "normalize_community_weight": True,
    "max_tokens": 12_000,  
    "context_name": "Reports",
}

map_llm_params = {
    "max_tokens": 1000,
    "temperature": 0.0,
    "response_format": {"type": "json_object"},
}

reduce_llm_params = {
    "max_tokens": 2000, 
    "temperature": 0.0,
}


search_engine = GlobalSearch(
    llm=llm,
    context_builder=context_builder,
    token_encoder=token_encoder,
    max_data_tokens=12_000,  
    map_llm_params=map_llm_params,
    reduce_llm_params=reduce_llm_params,
    allow_general_knowledge=False, 
    json_mode=True,  
    context_builder_params=context_builder_params,
    concurrent_coroutines=32,
    response_type="multiple paragraphs",  
)

result = await search_engine.asearch("贾府有多少人")
print(result.response)

# 贾府家族成员的人数众多，具体人数无法精确统计。根据不同分析师的报告，贾府家族成员至少有十几人，
# 也可能有数十人。这些成员包括贾政、王氏、贾珠、小姐、宁国公、贾代化、贾敬、惜春、贾珍、
# 贾蓉、贾雨村、甄士隐、林如海、黛玉等，以及贾化、贾琏、贾环、贾赦、贾敏等人。
# 家族成员之间的关系复杂，涉及父子、夫妻、兄弟姐妹、亲戚等，且可能涉及官位继承、
# 家族荣誉和家庭关系，因此具体人数可能更多 
# [Data: Reports (12, 27, 7, 14, 21, 6, 13, 26, 3, 30, 18, 23, 29, 22, 34), (33, 31, 35, 28, 32)]。

print(
    f"LLM calls: {result.llm_calls}. Prompt tokens: {result.prompt_tokens}. Output tokens: {result.output_tokens}."
)

## LLM calls: 3. Prompt tokens: 25247. Output tokens: 1299.

import os
import asyncio
from lightrag import LightRAG, QueryParam
from lightrag.llm import openai_complete_if_cache, openai_embedding, ollama_embedding
from lightrag.utils import EmbeddingFunc
import numpy as np

WORKING_DIR = "./dickens"

if not os.path.exists(WORKING_DIR):
    os.mkdir(WORKING_DIR)


async def llm_model_func(
    prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
) -> str:
    return await openai_complete_if_cache(
        "glm-4-flash",
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        api_key="**********************",
        base_url="https://open.bigmodel.cn/api/paas/v4/",
        **kwargs,
    )


async def embedding_func(texts: list[str]) -> np.ndarray:
    return await ollama_embedding(
        texts,
        embed_model="nomic-embed-text",
        host="http://localhost:11434",
    )


async def get_embedding_dim():
    test_text = ["This is a test sentence."]
    embedding = await embedding_func(test_text)
    # embedding_dim = embedding.shape[1]
    embedding_dim = len(embedding[0])
    print(f"Detected embedding dimension: {embedding_dim}")

    return embedding_dim


# function test
async def test_funcs():
    result = await llm_model_func("How are you?")
    print("llm_model_func: ", result)

    result = await embedding_func(["How are you?"])
    print("embedding_func: ", result)


# await test_funcs() # 测试llm+Embedding模型
# await get_embedding_dim()



embedding_dimension = await get_embedding_dim()
print(f"Detected embedding dimension: {embedding_dimension}")

rag = LightRAG(
    working_dir=WORKING_DIR,
    llm_model_func=llm_model_func,
    embedding_func=EmbeddingFunc(
        embedding_dim=embedding_dimension,
        max_token_size=8192,
        func=lambda texts: embedding_func(texts),
    ),
    addon_params = {"language": "Chinese"}
)

with open("./input/《红楼梦》 第一回 甄士隐梦幻识通灵 贾雨村风尘怀闺秀.txt", "r", encoding="utf-8") as f:
    text1 = f.read()
with open("./input/《红楼梦》 第二回 贾夫人仙逝扬州城　冷子兴演说荣国府.txt", "r", encoding="utf-8") as f:
    text2 = f.read()
all_text = [text1,text2]

await rag.ainsert(all_text)

from IPython.display import Markdown, display


# 五种查询模式
result = await rag.aquery("请介绍下贾母", param=QueryParam(mode="naive"))
display(Markdown(result))


result = await rag.aquery("请介绍下贾母", param=QueryParam(mode="local"))
display(Markdown(result))


result = await rag.aquery("请介绍下贾母", param=QueryParam(mode="global"))
display(Markdown(result))


result = await rag.aquery("请介绍下贾母", param=QueryParam(mode="hybrid"))
display(Markdown(result))

result = await rag.aquery("请介绍下贾母", param=QueryParam(mode="mix"))
display(Markdown(result))

Given the following description of a dataset:

{total_description}

Please identify 5 potential users who would engage with this dataset. For each user, list 5 tasks they would perform with this dataset. Then, for each (user, task) combination, generate 5 questions that require a high-level understanding of the entire dataset.

Output the results in the following structure:
- User 1: [user description]
    - Task 1: [task description]
        - Question 1:
        - Question 2:
        - Question 3:
        - Question 4:
        - Question 5:
    - Task 2: [task description]
        ...
    - Task 5: [task description]
- User 2: [user description]
    ...
- User 5: [user description]
    ...

Here are 5 potential users and their tasks with the dataset, along with 5 questions for each (user, task) combination:

**User 1: Beekeeper/Researcher**

- **Task 1: Monitor Bee Health**
    - Question 1: How does the dataset compare the prevalence of diseases like American Foulbrood and Chalkbrood across different regions and over time?
    - Question 2: Which factors (e.g., climate, forage availability, beekeeping practices) are most strongly associated with increased disease prevalence?
    - Question 3: How does the dataset compare the effectiveness of different treatments for Varroa mites across different regions?
    - Question 4: Are there any patterns in the timing of disease outbreaks that could be used to predict and prevent future outbreaks?
    - Question 5: How does the dataset compare the survival rates of honey bee colonies in different regions and over time?

- **Task 2: Analyze Forage Availability**
    - Question 1: How does the dataset quantify the availability of different forage sources (e.g., flowering plants, trees, shrubs) across different regions and over time?
    - Question 2: Which forage sources are most important for honey bee nutrition and colony health?
    - Question 3: How does the dataset compare the timing and duration of forage availability across different regions?
    - Question 4: Are there any patterns in the availability of forage sources that could be used to optimize beekeeping practices?
    - Question 5: How does the dataset compare the impact of different land management practices (e.g., agriculture, forestry, urban development) on forage availability?

- **Task 3: Study Bee Behavior**
    - Question 1: How does the dataset describe the foraging behavior of honey bees, including their preferred forage sources and foraging distances?
    - Question 2: Are there any patterns in the foraging behavior of honey bees that could be used to optimize beekeeping practices?
    - Question 3: How does the dataset compare the foraging behavior of different bee species?
    - Question 4: How does the dataset compare the foraging behavior of honey bees in different regions?
    - Question 5: How does the dataset compare the foraging behavior of honey bees over time?

- **Task 4: Analyze Beekeeping Practices**
    - Question 1: How does the dataset compare different beekeeping practices (e.g., hive management, feeding, disease control) across different regions?
    - Question 2: Which beekeeping practices are most effective for maintaining colony health and productivity?
    - Question 3: Are there any patterns in the effectiveness of different beekeeping practices that could be used to optimize beekeeping practices?
    - Question 4: How does the dataset compare the economic viability of different beekeeping practices?
    - Question 5: How does the dataset compare the sustainability of different beekeeping practices?

- **Task 5: Develop Predictive Models**
    - Question 1: Can the dataset be used to develop predictive models for honey bee colony health, disease prevalence, and forage availability?
    - Question 2: How accurate are the predictive models developed using the dataset?
    - Question 3: Can the predictive models be used to optimize beekeeping practices and improve honey bee health and productivity?
    - Question 4: Can the predictive models be used to identify potential threats to honey bee populations and develop mitigation strategies?
    - Question 5: How can the predictive models be improved and extended to other aspects of beekeeping and environmental monitoring?

---Role---
You are an expert tasked with evaluating two answers to the same question based on three criteria: **Comprehensiveness**, **Diversity**, and **Empowerment**.
---Goal---
You will evaluate two answers to the same question based on three criteria: **Comprehensiveness**, **Diversity**, and **Empowerment**.

- **Comprehensiveness**: How much detail does the answer provide to cover all aspects and details of the question?
- **Diversity**: How varied and rich is the answer in providing different perspectives and insights on the question?
- **Empowerment**: How well does the answer help the reader understand and make informed judgments about the topic?

For each criterion, choose the better answer (either Answer 1 or Answer 2) and explain why. Then, select an overall winner based on these three categories.

Here is the question:
{query}

Here are the two answers:

**Answer 1:**
{answer1}

**Answer 2:**
{answer2}

Evaluate both answers using the three criteria listed above and provide detailed explanations for each criterion.

Output your evaluation in the following JSON format:

{{
    "Comprehensiveness": {{
        "Winner": "[Answer 1 or Answer 2]",
        "Explanation": "[Provide explanation here]"
    }},
    "Empowerment": {{
        "Winner": "[Answer 1 or Answer 2]",
        "Explanation": "[Provide explanation here]"
    }},
    "Overall Winner": {{
        "Winner": "[Answer 1 or Answer 2]",
        "Explanation": "[Summarize why this answer is the overall winner based on the three criteria]"
    }}
}}