{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) 2024 Microsoft Corporation.\n",
    "# Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example of indexing from an existing in-memory dataframe\n",
    "\n",
    "Newer versions of GraphRAG let you submit a dataframe directly instead of running through the input processing step. This notebook demonstrates with regular or update runs.\n",
    "\n",
    "If performing an update, the assumption is that your dataframe contains only the new documents to add to the index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from pprint import pprint\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import graphrag.api as api\n",
    "from graphrag.config.load_config import load_config\n",
    "from graphrag.index.typing.pipeline_run_result import PipelineRunResult"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT_DIRECTORY = \"<your project directory>\"\n",
    "UPDATE = False\n",
    "FILENAME = \"new_documents.parquet\" if UPDATE else \"<original_documents>.parquet\"\n",
    "inputs = pd.read_parquet(f\"{PROJECT_DIRECTORY}/input/{FILENAME}\")\n",
    "# Only the bare minimum for input. These are the same fields that would be present after the load_input_documents workflow\n",
    "inputs = inputs.loc[:, [\"id\", \"title\", \"text\", \"creation_date\"]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate a `GraphRagConfig` object"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graphrag_config = load_config(Path(PROJECT_DIRECTORY))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Indexing API\n",
    "\n",
    "*Indexing* is the process of ingesting raw text data and constructing a knowledge graph. GraphRAG currently supports plaintext (`.txt`) and `.csv` file formats."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build an index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index_result: list[PipelineRunResult] = await api.build_index(\n",
    "    config=graphrag_config, input_documents=inputs, is_update_run=UPDATE\n",
    ")\n",
    "\n",
    "# index_result is a list of workflows that make up the indexing pipeline that was run\n",
    "for workflow_result in index_result:\n",
    "    status = f\"error\\n{workflow_result.errors}\" if workflow_result.errors else \"success\"\n",
    "    print(f\"Workflow Name: {workflow_result.workflow}\\tStatus: {status}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Query an index\n",
    "\n",
    "To query an index, several index files must first be read into memory and passed to the query API. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "entities = pd.read_parquet(f\"{PROJECT_DIRECTORY}/output/entities.parquet\")\n",
    "communities = pd.read_parquet(f\"{PROJECT_DIRECTORY}/output/communities.parquet\")\n",
    "community_reports = pd.read_parquet(\n",
    "    f\"{PROJECT_DIRECTORY}/output/community_reports.parquet\"\n",
    ")\n",
    "\n",
    "response, context = await api.global_search(\n",
    "    config=graphrag_config,\n",
    "    entities=entities,\n",
    "    communities=communities,\n",
    "    community_reports=community_reports,\n",
    "    community_level=2,\n",
    "    dynamic_community_selection=False,\n",
    "    response_type=\"Multiple Paragraphs\",\n",
    "    query=\"What are the top five themes of the dataset?\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The response object is the official reponse from graphrag while the context object holds various metadata regarding the querying process used to obtain the final response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Digging into the context a bit more provides users with extremely granular information such as what sources of data (down to the level of text chunks) were ultimately retrieved and used as part of the context sent to the LLM model)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(context)  # noqa: T203"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "graphrag",
   "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.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
