{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nCopyright (c) Microsoft Corporation.\\n'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Copyright (c) 2024 Microsoft Corporation.\n",
    "# Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "import tiktoken\n",
    "\n",
    "from graphrag.query.indexer_adapters import read_indexer_entities, read_indexer_reports\n",
    "from graphrag.query.llm.oai.chat_openai import ChatOpenAI\n",
    "from graphrag.query.llm.oai.typing import OpenaiApiType\n",
    "from graphrag.query.structured_search.global_search.community_context import (\n",
    "    GlobalCommunityContext,\n",
    ")\n",
    "from graphrag.query.structured_search.global_search.search import GlobalSearch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Global Search example\n",
    "\n",
    "Global search method generates answers by searching over all AI-generated community reports in a map-reduce fashion. This is a resource-intensive method, but often gives good responses for questions that require an understanding of the dataset as a whole (e.g. What are the most significant values of the herbs mentioned in this notebook?)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LLM setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "api_key = os.environ[\"GRAPHRAG_API_KEY\"]\n",
    "llm_model = os.environ[\"GRAPHRAG_LLM_MODEL\"]\n",
    "\n",
    "llm = ChatOpenAI(\n",
    "    api_key=api_key,\n",
    "    model=llm_model,\n",
    "    api_type=OpenaiApiType.OpenAI,  # OpenaiApiType.OpenAI or OpenaiApiType.AzureOpenAI\n",
    "    max_retries=20,\n",
    ")\n",
    "\n",
    "token_encoder = tiktoken.get_encoding(\"cl100k_base\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load community reports as context for global search\n",
    "\n",
    "- Load all community reports in the `create_final_community_reports` table from the ire-indexing engine, to be used as context data for global search.\n",
    "- Load entities from the `create_final_nodes` and `create_final_entities` tables from the ire-indexing engine, to be used for calculating community weights for context ranking. Note that this is optional (if no entities are provided, we will not calculate community weights and only use the `rank` attribute in the community reports table for context ranking)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parquet files generated from indexing pipeline\n",
    "INPUT_DIR = \"./inputs/operation dulce\"\n",
    "COMMUNITY_REPORT_TABLE = \"create_final_community_reports\"\n",
    "ENTITY_TABLE = \"create_final_nodes\"\n",
    "ENTITY_EMBEDDING_TABLE = \"create_final_entities\"\n",
    "\n",
    "# community level in the Leiden community hierarchy from which we will load the community reports\n",
    "# higher value means we use reports from more fine-grained communities (at the cost of higher computation cost)\n",
    "COMMUNITY_LEVEL = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "entity_df = pd.read_parquet(f\"{INPUT_DIR}/{ENTITY_TABLE}.parquet\")\n",
    "report_df = pd.read_parquet(f\"{INPUT_DIR}/{COMMUNITY_REPORT_TABLE}.parquet\")\n",
    "entity_embedding_df = pd.read_parquet(f\"{INPUT_DIR}/{ENTITY_EMBEDDING_TABLE}.parquet\")\n",
    "\n",
    "reports = read_indexer_reports(report_df, entity_df, COMMUNITY_LEVEL)\n",
    "entities = read_indexer_entities(entity_df, entity_embedding_df, COMMUNITY_LEVEL)\n",
    "print(f\"Report records: {len(report_df)}\")\n",
    "report_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Build global context based on community reports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "context_builder = GlobalCommunityContext(\n",
    "    community_reports=reports,\n",
    "    entities=entities,  # default to None if you don't want to use community weights for ranking\n",
    "    token_encoder=token_encoder,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Perform global search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "context_builder_params = {\n",
    "    \"use_community_summary\": False,  # False means using full community reports. True means using community short summaries.\n",
    "    \"shuffle_data\": True,\n",
    "    \"include_community_rank\": True,\n",
    "    \"min_community_rank\": 0,\n",
    "    \"community_rank_name\": \"rank\",\n",
    "    \"include_community_weight\": True,\n",
    "    \"community_weight_name\": \"occurrence weight\",\n",
    "    \"normalize_community_weight\": True,\n",
    "    \"max_tokens\": 12_000,  # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000)\n",
    "    \"context_name\": \"Reports\",\n",
    "}\n",
    "\n",
    "map_llm_params = {\n",
    "    \"max_tokens\": 1000,\n",
    "    \"temperature\": 0.0,\n",
    "    \"response_format\": {\"type\": \"json_object\"},\n",
    "}\n",
    "\n",
    "reduce_llm_params = {\n",
    "    \"max_tokens\": 2000,  # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 1000-1500)\n",
    "    \"temperature\": 0.0,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_engine = GlobalSearch(\n",
    "    llm=llm,\n",
    "    context_builder=context_builder,\n",
    "    token_encoder=token_encoder,\n",
    "    max_data_tokens=12_000,  # change this based on the token limit you have on your model (if you are using a model with 8k limit, a good setting could be 5000)\n",
    "    map_llm_params=map_llm_params,\n",
    "    reduce_llm_params=reduce_llm_params,\n",
    "    allow_general_knowledge=False,  # set this to True will add instruction to encourage the LLM to incorporate general knowledge in the response, which may increase hallucinations, but could be useful in some use cases.\n",
    "    json_mode=True,  # set this to False if your LLM model does not support JSON mode.\n",
    "    context_builder_params=context_builder_params,\n",
    "    concurrent_coroutines=32,\n",
    "    response_type=\"multiple paragraphs\",  # free form text describing the response type and format, can be anything, e.g. prioritized list, single paragraph, multiple paragraphs, multiple-page report\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = await search_engine.asearch(\n",
    "    \"What is the major conflict in this story and who are the protagonist and antagonist?\"\n",
    ")\n",
    "\n",
    "print(result.response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# inspect the data used to build the context for the LLM responses\n",
    "result.context_data[\"reports\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLM calls: 13. LLM tokens: 184660\n"
     ]
    }
   ],
   "source": [
    "# inspect number of LLM calls and tokens\n",
    "print(f\"LLM calls: {result.llm_calls}. LLM tokens: {result.prompt_tokens}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}