chore: import upstream snapshot with attribution

docs/query/drift_search.md

@@ -0,0 +1,36 @@
+# DRIFT Search 🔎
+
+## Combining Local and Global Search
+
+GraphRAG is a technique that uses large language models (LLMs) to create knowledge graphs and summaries from unstructured text documents and leverages them to improve retrieval-augmented generation (RAG) operations on private datasets. It offers comprehensive global overviews of large, private troves of unstructured text documents while also enabling exploration of detailed, localized information. By using LLMs to create comprehensive knowledge graphs that connect and describe entities and relationships contained in those documents, GraphRAG leverages semantic structuring of the data to generate responses to a wide variety of complex user queries.
+
+DRIFT search (Dynamic Reasoning and Inference with Flexible Traversal) builds upon Microsoft’s GraphRAG technique, combining characteristics of both global and local search to generate detailed responses in a method that balances computational costs with quality outcomes using our [drift search](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/drift_search/) method.
+
+## Methodology
+
+<p align="center">
+<img src="../../img/drift-search-diagram.png" alt="Figure 1. An entire DRIFT search hierarchy highlighting the three core phases of the DRIFT search process." align="center" />
+</p>
+<p align="center"><i><small>
+Figure 1. An entire DRIFT search hierarchy highlighting the three core phases of the DRIFT search process. A (Primer): DRIFT compares the user’s query with the top K most semantically relevant community reports, generating a broad initial answer and follow-up questions to steer further exploration. B (Follow-Up): DRIFT uses local search to refine queries, producing additional intermediate answers and follow-up questions that enhance specificity, guiding the engine towards context-rich information. A glyph on each node in the diagram shows the confidence the algorithm has to continue the query expansion step.  C (Output Hierarchy): The final output is a hierarchical structure of questions and answers ranked by relevance, reflecting a balanced mix of global insights and local refinements, making the results adaptable and comprehensive.</small></i></p>
+
+
+DRIFT Search introduces a new approach to local search queries by including community information in the search process. This greatly expands the breadth of the query’s starting point and leads to retrieval and usage of a far higher variety of facts in the final answer. This addition expands the GraphRAG query engine by providing a more comprehensive option for local search, which uses community insights to refine a query into detailed follow-up questions.
+
+## Configuration
+
+Below are the key parameters of the [DRIFTSearch class](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/drift_search/search.py):
+
+* `model`: Language model chat completion object to be used for response generation
+- `context_builder`: [context builder](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/drift_search/drift_context.py) object to be used for preparing context data from community reports and query information
+- `config`: model to define the DRIFT Search hyperparameters. [DRIFT Config model](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/config/models/drift_search_config.py)
+- `tokenizer`: token encoder for tracking the budget for the algorithm.
+- `query_state`: a state object as defined in [Query State](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/drift_search/state.py) that allows to track execution of a DRIFT Search instance, alongside follow ups and [DRIFT actions](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/drift_search/action.py).
+
+## How to Use
+
+An example of a drift search scenario can be found in the following [notebook](../examples_notebooks/drift_search.ipynb).
+
+## Learn More
+
+For a more in-depth look at the DRIFT search method, please refer to our [DRIFT Search blog post](https://www.microsoft.com/en-us/research/blog/introducing-drift-search-combining-global-and-local-search-methods-to-improve-quality-and-efficiency/)

docs/query/global_search.md

@@ -0,0 +1,73 @@
+# Global Search 🔎
+
+## Whole Dataset Reasoning
+
+Baseline RAG struggles with queries that require aggregation of information across the dataset to compose an answer. Queries such as “What are the top 5 themes in the data?” perform terribly because baseline RAG relies on a vector search of semantically similar text content within the dataset. There is nothing in the query to direct it to the correct information.
+
+However, with GraphRAG we can answer such questions, because the structure of the LLM-generated knowledge graph tells us about the structure (and thus themes) of the dataset as a whole. This allows the private dataset to be organized into meaningful semantic clusters that are pre-summarized. Using our [global search](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/global_search/) method, the LLM uses these clusters to summarize these themes when responding to a user query.
+
+## Methodology
+
+```mermaid
+---
+title: Global Search Dataflow
+---
+%%{ init: { 'flowchart': { 'curve': 'step' } } }%%
+flowchart LR
+
+    uq[User Query] --- .1
+    ch1[Conversation History] --- .1
+
+    subgraph RIR
+        direction TB
+        ri1[Rated Intermediate<br/>Response 1]~~~ri2[Rated Intermediate<br/>Response 2] -."{1..N}".-rin[Rated Intermediate<br/>Response N]
+    end
+
+    .1--Shuffled Community<br/>Report Batch 1-->RIR
+    .1--Shuffled Community<br/>Report Batch 2-->RIR---.2
+    .1--Shuffled Community<br/>Report Batch N-->RIR
+
+    .2--Ranking +<br/>Filtering-->agr[Aggregated Intermediate<br/>Responses]-->res[Response]
+
+
+
+     classDef green fill:#26B653,stroke:#333,stroke-width:2px,color:#fff;
+     classDef turquoise fill:#19CCD3,stroke:#333,stroke-width:2px,color:#fff;
+     classDef rose fill:#DD8694,stroke:#333,stroke-width:2px,color:#fff;
+     classDef orange fill:#F19914,stroke:#333,stroke-width:2px,color:#fff;
+     classDef purple fill:#B356CD,stroke:#333,stroke-width:2px,color:#fff;
+     classDef invisible fill:#fff,stroke:#fff,stroke-width:0px,color:#fff, width:0px;
+     class uq,ch1 turquoise;
+     class ri1,ri2,rin rose;
+     class agr orange;
+     class res purple;
+     class .1,.2 invisible;
+
+```
+
+Given a user query and, optionally, the conversation history, the global search method uses a collection of LLM-generated community reports from a specified level of the graph's community hierarchy as context data to generate response in a map-reduce manner. At the `map` step, community reports are segmented into text chunks of pre-defined size. Each text chunk is then used to produce an intermediate response containing a list of point, each of which is accompanied by a numerical rating indicating the importance of the point. At the `reduce` step, a filtered set of the most important points from the intermediate responses are aggregated and used as the context to generate the final response. 
+
+The quality of the global search’s response can be heavily influenced by the level of the community hierarchy chosen for sourcing community reports. Lower hierarchy levels, with their detailed reports, tend to yield more thorough responses, but may also increase the time and LLM resources needed to generate the final response due to the volume of reports.
+
+
+## Configuration
+
+Below are the key parameters of the [GlobalSearch class](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/global_search/search.py):
+
+* `model`: Language model chat completion object to be used for response generation
+* `context_builder`: [context builder](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/global_search/community_context.py) object to be used for preparing context data from community reports
+* `map_system_prompt`: prompt template used in the `map` stage. Default template can be found at [map_system_prompt](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/prompts/query/global_search_map_system_prompt.py)
+* `reduce_system_prompt`: prompt template used in the `reduce` stage, default template can be found at [reduce_system_prompt](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/prompts/query/global_search_reduce_system_prompt.py)
+* `response_type`: free-form text describing the desired response type and format (e.g., `Multiple Paragraphs`, `Multi-Page Report`)
+* `allow_general_knowledge`: setting this to True will include additional instructions to the `reduce_system_prompt` to prompt the LLM to incorporate relevant real-world knowledge outside of the dataset. Note that this may increase hallucinations, but can be useful for certain scenarios. Default is False
+*`general_knowledge_inclusion_prompt`: instruction to add to the `reduce_system_prompt` if `allow_general_knowledge` is enabled. Default instruction can be found at [general_knowledge_instruction](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/prompts/query/global_search_knowledge_system_prompt.py)
+* `max_data_tokens`: token budget for the context data
+* `map_llm_params`: a dictionary of additional parameters (e.g., temperature, max_tokens) to be passed to the LLM call at the `map` stage
+* `reduce_llm_params`: a dictionary of additional parameters (e.g., temperature, max_tokens) to passed to the LLM call at the `reduce` stage
+* `context_builder_params`: a dictionary of additional parameters to be passed to the [`context_builder`](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/global_search/community_context.py) object when building context window for the `map` stage.
+* `concurrent_coroutines`: controls the degree of parallelism in the `map` stage.
+* `callbacks`: optional callback functions, can be used to provide custom event handlers for LLM's completion streaming events
+
+## How to Use
+
+An example of a global search scenario can be found in the following [notebook](../examples_notebooks/global_search.ipynb).

docs/query/local_search.md

@@ -0,0 +1,62 @@
+# Local Search 🔎
+
+## Entity-based Reasoning
+
+The [local search](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/local_search/) method combines structured data from the knowledge graph with unstructured data from the input documents to augment the LLM context with relevant entity information at query time. It is well-suited for answering questions that require an understanding of specific entities mentioned in the input documents (e.g., “What are the healing properties of chamomile?”).
+
+## Methodology
+
+```mermaid
+---
+title: Local Search Dataflow
+---
+%%{ init: { 'flowchart': { 'curve': 'step' } } }%%
+flowchart LR
+
+    uq[User Query] ---.1
+    ch1[Conversation<br/>History]---.1
+
+    .1--Entity<br/>Description<br/>Embedding--> ee[Extracted Entities]
+
+    ee[Extracted Entities] ---.2--Entity-Text<br/>Unit Mapping--> ctu[Candidate<br/>Text Units]--Ranking + <br/>Filtering -->ptu[Prioritized<br/>Text Units]---.3
+    .2--Entity-Report<br/>Mapping--> ccr[Candidate<br/>Community Reports]--Ranking + <br/>Filtering -->pcr[Prioritized<br/>Community Reports]---.3
+    .2--Entity-Entity<br/>Relationships--> ce[Candidate<br/>Entities]--Ranking + <br/>Filtering -->pe[Prioritized<br/>Entities]---.3
+    .2--Entity-Entity<br/>Relationships--> cr[Candidate<br/>Relationships]--Ranking + <br/>Filtering -->pr[Prioritized<br/>Relationships]---.3
+    .2--Entity-Covariate<br/>Mappings--> cc[Candidate<br/>Covariates]--Ranking + <br/>Filtering -->pc[Prioritized<br/>Covariates]---.3
+    ch1 -->ch2[Conversation History]---.3
+    .3-->res[Response]
+
+     classDef green fill:#26B653,stroke:#333,stroke-width:2px,color:#fff;
+     classDef turquoise fill:#19CCD3,stroke:#333,stroke-width:2px,color:#fff;
+     classDef rose fill:#DD8694,stroke:#333,stroke-width:2px,color:#fff;
+     classDef orange fill:#F19914,stroke:#333,stroke-width:2px,color:#fff;
+     classDef purple fill:#B356CD,stroke:#333,stroke-width:2px,color:#fff;
+     classDef invisible fill:#fff,stroke:#fff,stroke-width:0px,color:#fff, width:0px;
+     class uq,ch1 turquoise
+     class ee green
+     class ctu,ccr,ce,cr,cc rose
+     class ptu,pcr,pe,pr,pc,ch2 orange
+     class res purple
+     class .1,.2,.3 invisible
+
+
+```
+
+Given a user query and, optionally, the conversation history, the local search method identifies a set of entities from the knowledge graph that are semantically-related to the user input. These entities serve as access points into the knowledge graph, enabling the extraction of further relevant details such as connected entities, relationships, entity covariates, and community reports. Additionally, it also extracts relevant text chunks from the raw input documents that are associated with the identified entities. These candidate data sources are then prioritized and filtered to fit within a single context window of pre-defined size, which is used to generate a response to the user query.
+
+## Configuration
+
+Below are the key parameters of the [LocalSearch class](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/local_search/search.py):
+
+* `model`: Language model chat completion object to be used for response generation
+* `context_builder`: [context builder](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/local_search/mixed_context.py) object to be used for preparing context data from collections of knowledge model objects
+* `system_prompt`: prompt template used to generate the search response. Default template can be found at [system_prompt](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/prompts/query/local_search_system_prompt.py)
+* `response_type`: free-form text describing the desired response type and format (e.g., `Multiple Paragraphs`, `Multi-Page Report`)
+* `llm_params`: a dictionary of additional parameters (e.g., temperature, max_tokens) to be passed to the LLM call
+* `context_builder_params`: a dictionary of additional parameters to be passed to the [`context_builder`](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/local_search/mixed_context.py) object when building context for the search prompt
+* `callbacks`: optional callback functions, can be used to provide custom event handlers for LLM's completion streaming events
+
+## How to Use
+
+An example of a local search scenario can be found in the following [notebook](../examples_notebooks/local_search.ipynb).
+

docs/query/notebooks/overview.md

@@ -0,0 +1,14 @@
+# API Notebooks
+
+- [API Overview Notebook](../../examples_notebooks/api_overview.ipynb)
+- [Bring-Your-Own Vector Store](../../examples_notebooks/custom_vector_store.ipynb)
+
+# Query Engine Notebooks
+
+For examples about running Query please refer to the following notebooks:
+
+- [Global Search Notebook](../../examples_notebooks/global_search.ipynb)
+- [Local Search Notebook](../../examples_notebooks/local_search.ipynb)
+- [DRIFT Search Notebook](../../examples_notebooks/drift_search.ipynb)
+
+The test dataset for these notebooks can be found in [dataset.zip](../../data/operation_dulce/dataset.zip){:download}.

docs/query/overview.md

@@ -0,0 +1,38 @@
+# Query Engine 🔎
+
+The Query Engine is the retrieval module of the GraphRAG library, and operates over completed [indexes](../index/overview.md).
+It is responsible for the following tasks:
+
+- [Local Search](#local-search)
+- [Global Search](#global-search)
+- [DRIFT Search](#drift-search)
+- Basic Search
+- [Question Generation](#question-generation)
+
+## Local Search
+
+Local search generates answers by combining relevant data from the AI-extracted knowledge-graph with text chunks of the raw documents. This method is suitable for questions that require an understanding of specific entities mentioned in the documents (e.g. What are the healing properties of chamomile?).
+
+For more details about how Local Search works please refer to the [Local Search](local_search.md) page.
+
+## Global Search
+
+Global search 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?).
+
+More about this is provided on the [Global Search](global_search.md) page.
+
+## DRIFT Search
+
+DRIFT Search introduces a new approach to local search queries by including community information in the search process. This greatly expands the breadth of the query’s starting point and leads to retrieval and usage of a far higher variety of facts in the final answer. This expands the GraphRAG query engine by providing a more comprehensive option for local search, which uses community insights to refine a query into detailed follow-up questions.
+
+To learn more about DRIFT Search, please refer to the [DRIFT Search](drift_search.md) page.
+
+## Basic Search
+
+GraphRAG includes a rudimentary implementation of basic vector RAG to make it easy to compare different search results based on the type of question you are asking. You can specify the top `k` text unit chunks to include in the summarization context.
+
+## Question Generation
+
+This functionality takes a list of user queries and generates the next candidate questions. This is useful for generating follow-up questions in a conversation or for generating a list of questions for the investigator to dive deeper into the dataset.
+
+Information about how question generation works can be found at the [Question Generation](question_generation.md) documentation page.

docs/query/question_generation.md

@@ -0,0 +1,23 @@
+# Question Generation ❔
+
+## Entity-based Question Generation
+
+The [question generation](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/question_gen/) method combines structured data from the knowledge graph with unstructured data from the input documents to generate candidate questions related to specific entities.
+
+## Methodology
+Given a list of prior user questions, the question generation method uses the same context-building approach employed in [local search](local_search.md) to extract and prioritize relevant structured and unstructured data, including entities, relationships, covariates, community reports and raw text chunks. These data records are then fitted into a single LLM prompt to generate candidate follow-up questions that represent the most important or urgent information content or themes in the data.
+
+## Configuration
+
+Below are the key parameters of the [Question Generation class](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/question_gen/local_gen.py):
+
+* `model`: Language model chat completion object to be used for response generation
+* `context_builder`: [context builder](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/local_search/mixed_context.py) object to be used for preparing context data from collections of knowledge model objects, using the same context builder class as in local search
+* `system_prompt`: prompt template used to generate candidate questions. Default template can be found at [system_prompt](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/prompts/query/question_gen_system_prompt.py)
+* `llm_params`: a dictionary of additional parameters (e.g., temperature, max_tokens) to be passed to the LLM call
+* `context_builder_params`: a dictionary of additional parameters to be passed to the [`context_builder`](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/query/structured_search/local_search/mixed_context.py) object when building context for the question generation prompt
+* `callbacks`: optional callback functions, can be used to provide custom event handlers for LLM's completion streaming events
+
+## How to Use
+
+An example of the question generation function can be found in the following [notebook](../examples_notebooks/local_search.ipynb).