diff --git a/.semversioner/next-release/patch-20260824180455220107.json b/.semversioner/next-release/patch-20260824180455220107.json
new file mode 100644
index 0000000000..a47d1ec47b
--- /dev/null
+++ b/.semversioner/next-release/patch-20260824180455220107.json
@@ -0,0 +1,4 @@
+{
+ "type": "patch",
+ "description": "Fix broken doc links and doc string parameter mismatches."
+}
diff --git a/docs/config/models.md b/docs/config/models.md
index 0bcdc2421a..401cba0318 100644
--- a/docs/config/models.md
+++ b/docs/config/models.md
@@ -37,7 +37,7 @@ Versions of GraphRAG before 2.2.0 made extensive use of `max_tokens` and `logit_
- Previously, GraphRAG used `max_tokens` to limit responses in a few locations. This is done so that we can have predictable content sizes when building downstream context windows for summarization. We have now switched from using `max_tokens` to use a prompted approach, which is working well in our tests. We suggest using `max_tokens` in your language model config only for budgetary reasons if you want to limit consumption, and not for expected response length control. We now also support the o-series equivalent `max_completion_tokens`, but if you use this keep in mind that there may be some unknown fixed reasoning consumption amount in addition to the response tokens, so it is not a good technique for response control.
- Previously, GraphRAG used a combination of `max_tokens` and `logit_bias` to strictly control a binary yes/no question during gleanings. This is not possible with reasoning models, so again we have switched to a prompted approach. Our tests with gpt-4o, gpt-4o-mini, and o1 show that this works consistently, but could have issues if you have an older or smaller model.
- The o-series models are much slower and more expensive. It may be useful to use an asymmetric approach to model use in your config: you can define as many models as you like in the `models` block of your settings.yaml and reference them by key for every workflow that requires a language model. You could use gpt-4o for indexing and o1 for query, for example. Experiment to find the right balance of cost, speed, and quality for your use case.
-- The o-series models contain a form of native native chain-of-thought reasoning that is absent in the non-o-series models. GraphRAG's prompts sometimes contain CoT because it was an effective technique with the gpt-4\* series. It may be counterproductive with the o-series, so you may want to tune or even re-write large portions of the prompt templates (particularly for graph and claim extraction).
+- The o-series models contain a form of native chain-of-thought reasoning that is absent in the non-o-series models. GraphRAG's prompts sometimes contain CoT because it was an effective technique with the gpt-4\* series. It may be counterproductive with the o-series, so you may want to tune or even re-write large portions of the prompt templates (particularly for graph and claim extraction).
Example config with asymmetric model use:
diff --git a/docs/config/yaml.md b/docs/config/yaml.md
index 453eac831e..af9e00c273 100644
--- a/docs/config/yaml.md
+++ b/docs/config/yaml.md
@@ -356,7 +356,7 @@ These are the settings used for Leiden hierarchical clustering of the graph to c
- `data_max_tokens` **int** - The maximum tokens to use constructing the final response from the reduces responses.
- `map_max_length` **int** - The maximum length to request for map responses, in words.
- `reduce_max_length` **int** - The maximum length to request for reduce responses, in words.
-- `dynamic_search_threshold` **int** - Rating threshold in include a community report.
+- `dynamic_search_threshold` **int** - Rating threshold to include a community report.
- `dynamic_search_keep_parent` **bool** - Keep parent community if any of the child communities are relevant.
- `dynamic_search_num_repeats` **int** - Number of times to rate the same community report.
- `dynamic_search_use_summary` **bool** - Use community summary instead of full_context.
diff --git a/docs/index/architecture.md b/docs/index/architecture.md
index c5d86938fe..bf5f509fc7 100644
--- a/docs/index/architecture.md
+++ b/docs/index/architecture.md
@@ -1,4 +1,4 @@
-# Indexing Architecture
+# Indexing Architecture
## Key Concepts
@@ -48,6 +48,6 @@ The following subsystems use a factory pattern that allows you to register your
- [vector store](https://github.com/microsoft/graphrag/blob/main/packages/graphrag-vectors/graphrag_vectors/vector_store_factory.py) - implement your own vector store other than the built-in lancedb, Azure AI Search, and CosmosDB ones built in
- [pipeline + workflows](https://github.com/microsoft/graphrag/blob/main/packages/graphrag/graphrag/index/workflows/factory.py) - implement your own workflow steps with a custom `run_workflow` function, or register an entire pipeline (list of named workflows)
-The links for each of these subsystems point to the source code of the factory, which includes registration of the default built-in implementations. In addition, we have a detailed discussion of [language models](../config/models.md), which includes and example of a custom provider, and a [sample notebook](../examples_notebooks/custom_vector_store.ipynb) that demonstrates a custom vector store.
+The links for each of these subsystems point to the source code of the factory, which includes registration of the default built-in implementations. In addition, we have a detailed discussion of [language models](../config/models.md), which includes an example of a custom provider, and a [sample notebook](../../packages/graphrag-vectors/example_notebooks/custom_vector_store.ipynb) that demonstrates a custom vector store.
-All of these factories allow you to register an impl using any string name you would like, even overriding built-in ones directly.
\ No newline at end of file
+All of these factories allow you to register an impl using any string name you would like, even overriding built-in ones directly.
diff --git a/docs/query/notebooks/overview.md b/docs/query/notebooks/overview.md
index 7df289e25f..f1162444e0 100644
--- a/docs/query/notebooks/overview.md
+++ b/docs/query/notebooks/overview.md
@@ -1,7 +1,7 @@
# API Notebooks
- [API Overview Notebook](../../examples_notebooks/api_overview.ipynb)
-- [Bring-Your-Own Vector Store](../../examples_notebooks/custom_vector_store.ipynb)
+- [Bring-Your-Own Vector Store](../../../packages/graphrag-vectors/example_notebooks/custom_vector_store.ipynb)
# Query Engine Notebooks
diff --git a/packages/graphrag-cache/README.md b/packages/graphrag-cache/README.md
index ac7709fc24..5054541cd2 100644
--- a/packages/graphrag-cache/README.md
+++ b/packages/graphrag-cache/README.md
@@ -1,22 +1,25 @@
# GraphRAG Cache
+> [!WARNING]
+> GraphRAG is a research project that explores the functional use of graphs to form a targeted context for question answering. Since our first release in July 2024 the capabilities of frontier models have changed dramatically, and our portfolio of research projects has diversified to match. This project is largely in maintenance mode, and won't be accepting new PRs or implementing new features. We'll perform bug fixes and dependency updates as appropriate, particularly to address CVEs as they arise.
+
This package contains a collection of utilities to handle GraphRAG caching implementation.
### Basic
-This example shows how to create a JSON cache with file storage using the GraphRAG cache package's configuration system.
+This example shows how to create a JSON cache with file storage using the GraphRAG cache package's configuration system.
-[Open the notebook to explore the basic example code](example-notebooks/basic_cache_example.ipynb)
+[Open the notebook to explore the basic example code](example_notebooks/basic_cache_example.ipynb)
### Custom Cache
This example demonstrates how to create a custom cache implementation by extending the base Cache class and registering it with the GraphRAG cache system. Once registered, the custom cache can be instantiated through the factory pattern using either CacheConfig or directly via cache_factory, allowing for extensible caching solutions tailored to specific needs.
-[Open the notebook to explore the basic custom example code](example-notebooks/custom_cache_example.ipynb)
+[Open the notebook to explore the basic custom example code](example_notebooks/custom_cache_example.ipynb)
#### Details
-By default, the `create_cache` comes with the following cache providers registered that correspond to the entries in the `CacheType` enum.
+By default, the `create_cache` comes with the following cache providers registered that correspond to the entries in the `CacheType` enum.
- `JsonCache`
- `MemoryCache`
@@ -40,4 +43,4 @@ cache = cache_factory.create(
)
...
-```
\ No newline at end of file
+```
diff --git a/packages/graphrag-chunking/README.md b/packages/graphrag-chunking/README.md
index 2415bbcfcd..6397fd13ea 100644
--- a/packages/graphrag-chunking/README.md
+++ b/packages/graphrag-chunking/README.md
@@ -1,5 +1,8 @@
# GraphRAG Chunking
+> [!WARNING]
+> GraphRAG is a research project that explores the functional use of graphs to form a targeted context for question answering. Since our first release in July 2024 the capabilities of frontier models have changed dramatically, and our portfolio of research projects has diversified to match. This project is largely in maintenance mode, and won't be accepting new PRs or implementing new features. We'll perform bug fixes and dependency updates as appropriate, particularly to address CVEs as they arise.
+
This package contains a collection of text chunkers, a core config model, and a factory for acquiring instances.
## Examples
@@ -16,7 +19,6 @@ The TokenChunker splits text into fixed-size chunks based on token count rather
[Open the notebook to explore the token chunking example code](example_notebooks/token_chunking_example.ipynb)
-
### Using the factory via helper util
The create_chunker factory function provides a configuration-driven approach to instantiate chunkers by accepting a ChunkingConfig object that specifies the chunking strategy and parameters. This allows for more flexible and maintainable code by separating chunker configuration from direct instantiation.
diff --git a/packages/graphrag-common/README.md b/packages/graphrag-common/README.md
index df114726c3..56e264f365 100644
--- a/packages/graphrag-common/README.md
+++ b/packages/graphrag-common/README.md
@@ -1,5 +1,8 @@
# GraphRAG Common
+> [!WARNING]
+> GraphRAG is a research project that explores the functional use of graphs to form a targeted context for question answering. Since our first release in July 2024 the capabilities of frontier models have changed dramatically, and our portfolio of research projects has diversified to match. This project is largely in maintenance mode, and won't be accepting new PRs or implementing new features. We'll perform bug fixes and dependency updates as appropriate, particularly to address CVEs as they arise.
+
This package provides utility modules for GraphRAG, including a flexible factory system for dependency injection and service registration, and a comprehensive configuration loading system with Pydantic model support, environment variable substitution, and automatic file discovery.
## Factory module
diff --git a/packages/graphrag-input/README.md b/packages/graphrag-input/README.md
index cae4cc3482..a19b6e04c4 100644
--- a/packages/graphrag-input/README.md
+++ b/packages/graphrag-input/README.md
@@ -1,5 +1,8 @@
# GraphRAG Inputs
+> [!WARNING]
+> GraphRAG is a research project that explores the functional use of graphs to form a targeted context for question answering. Since our first release in July 2024 the capabilities of frontier models have changed dramatically, and our portfolio of research projects has diversified to match. This project is largely in maintenance mode, and won't be accepting new PRs or implementing new features. We'll perform bug fixes and dependency updates as appropriate, particularly to address CVEs as they arise.
+
This package provides input document loading utilities for GraphRAG, supporting multiple file formats including CSV, JSON, JSON Lines, and plain text.
## Supported File Types
@@ -21,22 +24,21 @@ Additionally, we support the `InputType.MarkItDown` format, which uses the [Mark
Basic usage with the factory:
-
1. Import a pdf with MarkItDown:
```bash
pip install 'markitdown[pdf]' # required dependency for pdf processing
```
-
2. YAML config example for above:
+
```yaml
input:
type: markitdown
file_pattern: ".*\\.pdf$$"
input_storage:
- type: file
- base_dir: "input"
+ type: file
+ base_dir: "input"
```
-[Open the notebook to explore the input example code](example_notebooks/input_example.ipynb)
\ No newline at end of file
+[Open the notebook to explore the input example code](example_notebooks/input_example.ipynb)
diff --git a/packages/graphrag-llm/README.md b/packages/graphrag-llm/README.md
index 3c7b9af0d3..5e3a1f1e53 100644
--- a/packages/graphrag-llm/README.md
+++ b/packages/graphrag-llm/README.md
@@ -1,16 +1,18 @@
# GraphRAG LLM
+> [!WARNING]
+> GraphRAG is a research project that explores the functional use of graphs to form a targeted context for question answering. Since our first release in July 2024 the capabilities of frontier models have changed dramatically, and our portfolio of research projects has diversified to match. This project is largely in maintenance mode, and won't be accepting new PRs or implementing new features. We'll perform bug fixes and dependency updates as appropriate, particularly to address CVEs as they arise.
+
## Basic Completion
This example demonstrates basic usage of the LLM library to interact with Azure OpenAI. It loads environment variables for API configuration, creates a ModelConfig for Azure OpenAI, and sends a simple question to the model. The code handles both streaming and non-streaming responses (streaming responses are printed chunk by chunk in real-time, while non-streaming responses are printed all at once). It also shows how to use the gather_completion_response utility function as a simpler alternative that automatically handles both response types and returns the complete text.
[Open the notebook to explore the basic completion example code](example_notebooks/basic_completion_example.ipynb)
-
## Basic Embedding
This examples demonstrates how to generate text embeddings using the GraphRAG LLM library with Azure OpenAI's embedding service. It loads API credentials from environment variables, creates a ModelConfig for the Azure embedding model and configures authentication to use either API key or Azure Managed Identity. The script then creates an embedding client and processes a batch of two text strings ("Hello world" and "How are you?") to generate their vector embeddings.
-[Open the notebook to explore the basic embeddings example code](example_notebooks/basic_completion_example.ipynb)
+[Open the notebook to explore the basic embeddings example code](example_notebooks/basic_embedding_example.ipynb)
-View the [notebooks](example_notebooks/README.md) for more examples.
\ No newline at end of file
+View the [notebooks](notebooks/README.md) for more examples.
diff --git a/packages/graphrag-llm/graphrag_llm/utils/create_embedding_response.py b/packages/graphrag-llm/graphrag_llm/utils/create_embedding_response.py
index b651a32816..ed958246b4 100644
--- a/packages/graphrag-llm/graphrag_llm/utils/create_embedding_response.py
+++ b/packages/graphrag-llm/graphrag_llm/utils/create_embedding_response.py
@@ -13,9 +13,7 @@ def create_embedding_response(
Args:
embeddings: List of embedding vectors.
- model: The model used to create the embeddings.
-
- Returns
+ batch_size: The number of embedding objects to generate.
-------
An LLMEmbeddingResponse object.
"""
diff --git a/packages/graphrag-storage/README.md b/packages/graphrag-storage/README.md
index 5671d5e319..d673e2acc8 100644
--- a/packages/graphrag-storage/README.md
+++ b/packages/graphrag-storage/README.md
@@ -1,5 +1,8 @@
# GraphRAG Storage
+> [!WARNING]
+> GraphRAG is a research project that explores the functional use of graphs to form a targeted context for question answering. Since our first release in July 2024 the capabilities of frontier models have changed dramatically, and our portfolio of research projects has diversified to match. This project is largely in maintenance mode, and won't be accepting new PRs or implementing new features. We'll perform bug fixes and dependency updates as appropriate, particularly to address CVEs as they arise.
+
This package provides a unified storage abstraction layer with support for multiple backends including file system, Azure Blob, Azure Cosmos, and memory storage. It features a factory-based creation system with configuration-driven setup and extensible architecture for implementing custom storage providers.
## Basic
@@ -14,10 +17,9 @@ Here we create a custom storage implementation by extending the base Storage cla
[Open the notebook to explore the custom storage example code](example_notebooks/basic_storage_example.ipynb)
-
## Details
-By default, the `create_storage` comes with the following storage providers registered that correspond to the entries in the `StorageType` enum.
+By default, the `create_storage` comes with the following storage providers registered that correspond to the entries in the `StorageType` enum.
- `FileStorage`
- `AzureBlobStorage`
@@ -42,4 +44,4 @@ storage = storage_factory.create(
)
...
-```
\ No newline at end of file
+```
diff --git a/packages/graphrag-vectors/README.md b/packages/graphrag-vectors/README.md
index 082f33fcc7..3cab06629a 100644
--- a/packages/graphrag-vectors/README.md
+++ b/packages/graphrag-vectors/README.md
@@ -1,5 +1,8 @@
# GraphRAG Vectors
+> [!WARNING]
+> GraphRAG is a research project that explores the functional use of graphs to form a targeted context for question answering. Since our first release in July 2024 the capabilities of frontier models have changed dramatically, and our portfolio of research projects has diversified to match. This project is largely in maintenance mode, and won't be accepting new PRs or implementing new features. We'll perform bug fixes and dependency updates as appropriate, particularly to address CVEs as they arise.
+
This package provides vector store implementations for GraphRAG with support for multiple backends including LanceDB, Azure AI Search, and Azure Cosmos DB. It offers both a convenient configuration-driven API and direct factory access for creating and managing vector stores with flexible index schema definitions.
## Basic usage with the utility function (recommended)
@@ -29,5 +32,6 @@ You can register custom vector store implementations:
## Configuration
Vector stores are configured using:
+
- `VectorStoreConfig`: baseline parameters for the store
- `IndexSchema`: Schema configuration for the specific index to create/connect to (index name, field names, vector size)
diff --git a/packages/graphrag/README.md b/packages/graphrag/README.md
index 23748b176e..8693566f2a 100644
--- a/packages/graphrag/README.md
+++ b/packages/graphrag/README.md
@@ -1,5 +1,8 @@
# GraphRAG
+> [!WARNING]
+> GraphRAG is a research project that explores the functional use of graphs to form a targeted context for question answering. Since our first release in July 2024 the capabilities of frontier models have changed dramatically, and our portfolio of research projects has diversified to match. This project is largely in maintenance mode, and won't be accepting new PRs or implementing new features. We'll perform bug fixes and dependency updates as appropriate, particularly to address CVEs as they arise.
+
👉 [Microsoft Research Blog Post](https://www.microsoft.com/en-us/research/blog/graphrag-unlocking-llm-discovery-on-narrative-private-data/)
👉 [Read the docs](https://microsoft.github.io/graphrag)
👉 [GraphRAG Arxiv](https://arxiv.org/pdf/2404.16130)
@@ -33,12 +36,12 @@ To get started with the GraphRAG system we recommend trying the [command line qu
This repository presents a methodology for using knowledge graph memory structures to enhance LLM outputs. Please note that the provided code serves as a demonstration and is not an officially supported Microsoft offering.
-⚠️ *Warning: GraphRAG indexing can be an expensive operation, please read all of the documentation to understand the process and costs involved, and start small.*
+⚠️ _Warning: GraphRAG indexing can be an expensive operation, please read all of the documentation to understand the process and costs involved, and start small._
## Diving Deeper
-- To learn about our contribution guidelines, see [CONTRIBUTING.md](./CONTRIBUTING.md)
-- To start developing _GraphRAG_, see [DEVELOPING.md](./DEVELOPING.md)
+- To learn about our contribution guidelines, see [CONTRIBUTING.md](../../CONTRIBUTING.md)
+- To start developing _GraphRAG_, see [DEVELOPING.md](../../DEVELOPING.md)
- Join the conversation and provide feedback in the [GitHub Discussions tab!](https://github.com/microsoft/graphrag/discussions)
## Prompt Tuning
@@ -48,20 +51,20 @@ We strongly recommend to fine-tune your prompts following the [Prompt Tuning Gui
## Versioning
-Please see the [breaking changes](./breaking-changes.md) document for notes on our approach to versioning the project.
+Please see the [breaking changes](../../breaking-changes.md) document for notes on our approach to versioning the project.
-*Always run `uv run poe init --root [path] --force` between minor version bumps to ensure you have the latest config format. Run the provided migration notebook between major version bumps if you want to avoid re-indexing prior datasets. Note that this will overwrite your configuration and prompts, so back them up if necessary.*
+_Always run `uv run poe init --root [path] --force` between minor version bumps to ensure you have the latest config format. Run the provided migration notebook between major version bumps if you want to avoid re-indexing prior datasets. Note that this will overwrite your configuration and prompts, so back them up if necessary._
## Responsible AI FAQ
-See [RAI_TRANSPARENCY.md](./RAI_TRANSPARENCY.md)
+See [RAI_TRANSPARENCY.md](../../RAI_TRANSPARENCY.md)
-- [What is GraphRAG?](./RAI_TRANSPARENCY.md#what-is-graphrag)
-- [What can GraphRAG do?](./RAI_TRANSPARENCY.md#what-can-graphrag-do)
-- [What are GraphRAG’s intended use(s)?](./RAI_TRANSPARENCY.md#what-are-graphrags-intended-uses)
-- [How was GraphRAG evaluated? What metrics are used to measure performance?](./RAI_TRANSPARENCY.md#how-was-graphrag-evaluated-what-metrics-are-used-to-measure-performance)
-- [What are the limitations of GraphRAG? How can users minimize the impact of GraphRAG’s limitations when using the system?](./RAI_TRANSPARENCY.md#what-are-the-limitations-of-graphrag-how-can-users-minimize-the-impact-of-graphrags-limitations-when-using-the-system)
-- [What operational factors and settings allow for effective and responsible use of GraphRAG?](./RAI_TRANSPARENCY.md#what-operational-factors-and-settings-allow-for-effective-and-responsible-use-of-graphrag)
+- [What is GraphRAG?](../../RAI_TRANSPARENCY.md#what-is-graphrag)
+- [What can GraphRAG do?](../../RAI_TRANSPARENCY.md#what-can-graphrag-do)
+- [What are GraphRAG’s intended use(s)?](../../RAI_TRANSPARENCY.md#what-are-graphrags-intended-uses)
+- [How was GraphRAG evaluated? What metrics are used to measure performance?](../../RAI_TRANSPARENCY.md#how-was-graphrag-evaluated-what-metrics-are-used-to-measure-performance)
+- [What are the limitations of GraphRAG? How can users minimize the impact of GraphRAG’s limitations when using the system?](../../RAI_TRANSPARENCY.md#what-are-the-limitations-of-graphrag-how-can-users-minimize-the-impact-of-graphrags-limitations-when-using-the-system)
+- [What operational factors and settings allow for effective and responsible use of GraphRAG?](../../RAI_TRANSPARENCY.md#what-operational-factors-and-settings-allow-for-effective-and-responsible-use-of-graphrag)
## Trademarks
diff --git a/packages/graphrag/graphrag/api/index.py b/packages/graphrag/graphrag/api/index.py
index cc9d0956d5..c666ad5aa3 100644
--- a/packages/graphrag/graphrag/api/index.py
+++ b/packages/graphrag/graphrag/api/index.py
@@ -43,12 +43,14 @@ async def build_index(
The configuration.
method : IndexingMethod default=IndexingMethod.Standard
Styling of indexing to perform (full LLM, NLP + LLM, etc.).
- memory_profile : bool
- Whether to enable memory profiling.
+ is_update_run : bool default=False
+ Whether this run is an incremental update over an existing index.
callbacks : list[WorkflowCallbacks] | None default=None
A list of callbacks to register.
additional_context : dict[str, Any] | None default=None
Additional context to pass to the pipeline run. This can be accessed in the pipeline state under the 'additional_context' key.
+ verbose : bool default=False
+ Whether to enable verbose logging.
input_documents : pd.DataFrame | None default=None.
Override document loading and parsing and supply your own dataframe of documents to index.
diff --git a/packages/graphrag/graphrag/config/models/basic_search_config.py b/packages/graphrag/graphrag/config/models/basic_search_config.py
index 5b48b74d82..b133ae5af6 100644
--- a/packages/graphrag/graphrag/config/models/basic_search_config.py
+++ b/packages/graphrag/graphrag/config/models/basic_search_config.py
@@ -9,7 +9,7 @@
class BasicSearchConfig(BaseModel):
- """The default configuration section for Cache."""
+ """The default configuration section for Basic Search."""
prompt: str | None = Field(
description="The basic search prompt to use.",
diff --git a/packages/graphrag/graphrag/config/models/drift_search_config.py b/packages/graphrag/graphrag/config/models/drift_search_config.py
index ce77fc6b1f..2a0c7c4733 100644
--- a/packages/graphrag/graphrag/config/models/drift_search_config.py
+++ b/packages/graphrag/graphrag/config/models/drift_search_config.py
@@ -9,7 +9,7 @@
class DRIFTSearchConfig(BaseModel):
- """The default configuration section for Cache."""
+ """The default configuration section for DRIFT Search."""
prompt: str | None = Field(
description="The drift search prompt to use.",
diff --git a/packages/graphrag/graphrag/config/models/global_search_config.py b/packages/graphrag/graphrag/config/models/global_search_config.py
index 1b48500ed2..94decac5f7 100644
--- a/packages/graphrag/graphrag/config/models/global_search_config.py
+++ b/packages/graphrag/graphrag/config/models/global_search_config.py
@@ -9,7 +9,7 @@
class GlobalSearchConfig(BaseModel):
- """The default configuration section for Cache."""
+ """The default configuration section for Global Search."""
map_prompt: str | None = Field(
description="The global search mapper prompt to use.",
@@ -46,7 +46,7 @@ class GlobalSearchConfig(BaseModel):
# configurations for dynamic community selection
dynamic_search_threshold: int = Field(
- description="Rating threshold in include a community report",
+ description="Rating threshold to include a community report",
default=graphrag_config_defaults.global_search.dynamic_search_threshold,
)
dynamic_search_keep_parent: bool = Field(
diff --git a/packages/graphrag/graphrag/config/models/local_search_config.py b/packages/graphrag/graphrag/config/models/local_search_config.py
index c9adbe9098..ef2212b026 100644
--- a/packages/graphrag/graphrag/config/models/local_search_config.py
+++ b/packages/graphrag/graphrag/config/models/local_search_config.py
@@ -9,7 +9,7 @@
class LocalSearchConfig(BaseModel):
- """The default configuration section for Cache."""
+ """The default configuration section for Local Search."""
prompt: str | None = Field(
description="The local search prompt to use.",
diff --git a/packages/graphrag/graphrag/index/update/communities.py b/packages/graphrag/graphrag/index/update/communities.py
index 213db2bb15..e03d556485 100644
--- a/packages/graphrag/graphrag/index/update/communities.py
+++ b/packages/graphrag/graphrag/index/update/communities.py
@@ -23,8 +23,6 @@ def _update_and_merge_communities(
The old communities.
delta_communities : pd.DataFrame
The delta communities.
- community_id_mapping : dict
- The community id mapping.
Returns
-------
diff --git a/packages/graphrag/graphrag/prompt_tune/generator/community_report_rating.py b/packages/graphrag/graphrag/prompt_tune/generator/community_report_rating.py
index a4cae4c75f..91dfea3e10 100644
--- a/packages/graphrag/graphrag/prompt_tune/generator/community_report_rating.py
+++ b/packages/graphrag/graphrag/prompt_tune/generator/community_report_rating.py
@@ -16,13 +16,13 @@
async def generate_community_report_rating(
model: "LLMCompletion", domain: str, persona: str, docs: str | list[str]
) -> str:
- """Generate an LLM persona to use for GraphRAG prompts.
+ """Generate a community report rating description for GraphRAG prompts.
Parameters
----------
- model (LLMCompletion): The LLM to use for generation
- domain (str): The domain to generate a rating for
- - persona (str): The persona to generate a rating for for
+ - persona (str): The persona to generate a rating for
- docs (str | list[str]): Documents used to contextualize the rating
Returns
diff --git a/packages/graphrag/graphrag/query/structured_search/drift_search/primer.py b/packages/graphrag/graphrag/query/structured_search/drift_search/primer.py
index d153abc72a..c83f822111 100644
--- a/packages/graphrag/graphrag/query/structured_search/drift_search/primer.py
+++ b/packages/graphrag/graphrag/query/structured_search/drift_search/primer.py
@@ -5,6 +5,7 @@
import logging
import secrets
+import textwrap
import time
from typing import TYPE_CHECKING
@@ -57,7 +58,7 @@ def __init__(
Initialize the PrimerQueryProcessor.
Args:
- chat_llm (ChatOpenAI): The language model used to process the query.
+ chat_model (LLMCompletion): The language model used to process the query.
text_embedder (BaseTextEmbedding): The text embedding model.
reports (list[CommunityReport]): List of community reports.
tokenizer (Tokenizer, optional): Token encoder for token counting.
@@ -80,10 +81,18 @@ async def expand_query(self, query: str) -> tuple[str, dict[str, int]]:
"""
template = secrets.choice(self.reports).full_content # nosec S311
- prompt = f"""Create a hypothetical answer to the following query: {query}\n\n
- Format it to follow the structure of the template below:\n\n
- {template}\n"
- Ensure that the hypothetical answer does not reference new named entities that are not present in the original query."""
+ prompt = (
+ textwrap.dedent(f"""\
+ Create a hypothetical answer to the following query: {query}
+
+ Format it to follow the structure of the template below:
+
+ """)
+ + template
+ + textwrap.dedent("""
+
+ Ensure that the hypothetical answer does not reference new named entities that are not present in the original query.""")
+ )
model_response: LLMCompletionResponse = await self.chat_model.completion_async(
messages=prompt
@@ -134,7 +143,7 @@ def __init__(
Args:
config (DRIFTSearchConfig): Configuration settings for DRIFT search.
- chat_llm (ChatOpenAI): The language model used for searching.
+ chat_model (LLMCompletion): The language model used for searching.
tokenizer (Tokenizer, optional): Tokenizer for managing tokens.
"""
self.chat_model = chat_model
diff --git a/packages/graphrag/graphrag/query/structured_search/drift_search/search.py b/packages/graphrag/graphrag/query/structured_search/drift_search/search.py
index 070c5b5519..c595735d49 100644
--- a/packages/graphrag/graphrag/query/structured_search/drift_search/search.py
+++ b/packages/graphrag/graphrag/query/structured_search/drift_search/search.py
@@ -49,11 +49,11 @@ def __init__(
Initialize the DRIFTSearch class.
Args:
- llm (ChatOpenAI): The language model used for searching.
+ model (LLMCompletion): The language model used for searching.
context_builder (DRIFTSearchContextBuilder): Builder for search context.
- config (DRIFTSearchConfig, optional): Configuration settings for DRIFTSearch.
tokenizer (Tokenizer, optional): Token encoder for managing tokens.
query_state (QueryState, optional): State of the current search query.
+ callbacks (list[QueryCallbacks], optional): Callbacks to invoke during search.
"""
super().__init__(model, context_builder, tokenizer)