Refactor config (#1593)

* Refactor config

- Add new ModelConfig to represent LLM settings
    - Combines LLMParameters, ParallelizationParameters, encoding_model, and async_mode
- Add top level models config that is a list of available LLM ModelConfigs
- Remove LLMConfig inheritance and delete LLMConfig
    - Replace the inheritance with a model_id reference to the ModelConfig listed in the top level models config
- Remove all fallbacks and hydration logic from create_graphrag_config
    - This removes the automatic env variable overrides
- Support env variables within config files using Templating
    - This requires "$" to be escaped with extra "$" so ".*\\.txt$" becomes ".*\\.txt$$"
- Update init content to initialize new config file with the ModelConfig structure

* Use dict of ModelConfig instead of list

* Add model validations and unit tests

* Fix ruff checks

* Add semversioner change

* Fix unit tests

* validate root_dir in pydantic model

* Rename ModelConfig to LanguageModelConfig

* Rename ModelConfigMissingError to LanguageModelConfigMissingError

* Add validationg for unexpected API keys

* Allow skipping pydantic validation for testing/mocking purposes.

* Add default lm configs to verb tests

* smoke test

* remove config from flows to fix llm arg mapping

* Fix embedding llm arg mapping

* Remove timestamp from smoke test outputs

* Remove unused "subworkflows" smoke test properties

* Add models to smoke test configs

* Update smoke test output path

* Send logs to logs folder

* Fix output path

* Fix csv test file pattern

* Update placeholder

* Format

* Instantiate default model configs

* Fix unit tests for config defaults

* Fix migration notebook

* Remove create_pipeline_config

* Remove several unused config models

* Remove indexing embedding and input configs

* Move embeddings function to config

* Remove skip_workflows

* Remove skip embeddings in favor of explicit naming

* fix unit test spelling mistake

* self.models[model_id] is already a language model. Remove redundant casting.

* update validation errors to instruct users to rerun graphrag init

* instantiate LanguageModelConfigs with validation

* skip validation in unit tests

* update verb tests to use default model settings instead of skipping validation

* test using llm settings

* cleanup verb tests

* remove unsafe default model config

* remove the ability to skip pydantic validation

* remove None union types when default values are set

* move vector_store from embeddings to top level of config and delete resolve_paths

* update vector store settings

* fix vector store and smoke tests

* fix serializing vector_store settings

* fix vector_store usage

* fix vector_store type

* support cli overrides for loading graphrag config

* rename storage to output

* Add --force flag to init

* Remove run_id and resume, fix Drift config assignment

* Ruff

---------

Co-authored-by: Nathan Evans <[email protected]>
Co-authored-by: Alonso Guevara <[email protected]>

Author: 3 people

.semversioner/next-release/minor-20250109223356701278.json

@@ -0,0 +1,4 @@
+{
+  "type": "minor",
+  "description": "Remove config inheritance, hydration, and automatic env var overlays."
+}

docs/examples_notebooks/index_migration.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 66,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -25,7 +25,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 67,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -37,27 +37,28 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "from pathlib import Path\n",
     "\n",
     "from graphrag.config.load_config import load_config\n",
     "from graphrag.config.resolve_path import resolve_paths\n",
-    "from graphrag.index.create_pipeline_config import create_pipeline_config\n",
-    "from graphrag.storage.factory import create_storage\n",
+    "from graphrag.storage.factory import StorageFactory\n",
     "\n",
     "# This first block does some config loading, path resolution, and translation that is normally done by the CLI/API when running a full workflow\n",
     "config = load_config(Path(PROJECT_DIRECTORY))\n",
     "resolve_paths(config)\n",
-    "pipeline_config = create_pipeline_config(config)\n",
-    "storage = create_storage(pipeline_config.storage)"
+    "storage_config = config.storage.model_dump()  # type: ignore\n",
+    "storage = StorageFactory().create_storage(\n",
+    "    storage_type=storage_config[\"type\"], kwargs=storage_config\n",
+    ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 69,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -68,7 +69,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 63,
+   "execution_count": 70,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -97,7 +98,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 64,
+   "execution_count": 71,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -108,22 +109,16 @@
     "# First we'll go through any parquet files that had model changes and update them\n",
     "# The new data model may have removed excess columns as well, but we will only make the minimal changes required for compatibility\n",
     "\n",
-    "final_documents = await load_table_from_storage(\n",
-    "    \"create_final_documents.parquet\", storage\n",
-    ")\n",
-    "final_text_units = await load_table_from_storage(\n",
-    "    \"create_final_text_units.parquet\", storage\n",
-    ")\n",
-    "final_entities = await load_table_from_storage(\"create_final_entities.parquet\", storage)\n",
-    "final_nodes = await load_table_from_storage(\"create_final_nodes.parquet\", storage)\n",
+    "final_documents = await load_table_from_storage(\"create_final_documents\", storage)\n",
+    "final_text_units = await load_table_from_storage(\"create_final_text_units\", storage)\n",
+    "final_entities = await load_table_from_storage(\"create_final_entities\", storage)\n",
+    "final_nodes = await load_table_from_storage(\"create_final_nodes\", storage)\n",
     "final_relationships = await load_table_from_storage(\n",
-    "    \"create_final_relationships.parquet\", storage\n",
-    ")\n",
-    "final_communities = await load_table_from_storage(\n",
-    "    \"create_final_communities.parquet\", storage\n",
+    "    \"create_final_relationships\", storage\n",
     ")\n",
+    "final_communities = await load_table_from_storage(\"create_final_communities\", storage)\n",
     "final_community_reports = await load_table_from_storage(\n",
-    "    \"create_final_community_reports.parquet\", storage\n",
+    "    \"create_final_community_reports\", storage\n",
     ")\n",
     "\n",
     "\n",
@@ -183,44 +178,41 @@
     "        parent_df, on=\"community\", how=\"left\"\n",
     "    )\n",
     "\n",
-    "await write_table_to_storage(final_documents, \"create_final_documents.parquet\", storage)\n",
+    "await write_table_to_storage(final_documents, \"create_final_documents\", storage)\n",
+    "await write_table_to_storage(final_text_units, \"create_final_text_units\", storage)\n",
+    "await write_table_to_storage(final_entities, \"create_final_entities\", storage)\n",
+    "await write_table_to_storage(final_nodes, \"create_final_nodes\", storage)\n",
+    "await write_table_to_storage(final_relationships, \"create_final_relationships\", storage)\n",
+    "await write_table_to_storage(final_communities, \"create_final_communities\", storage)\n",
     "await write_table_to_storage(\n",
-    "    final_text_units, \"create_final_text_units.parquet\", storage\n",
-    ")\n",
-    "await write_table_to_storage(final_entities, \"create_final_entities.parquet\", storage)\n",
-    "await write_table_to_storage(final_nodes, \"create_final_nodes.parquet\", storage)\n",
-    "await write_table_to_storage(\n",
-    "    final_relationships, \"create_final_relationships.parquet\", storage\n",
-    ")\n",
-    "await write_table_to_storage(\n",
-    "    final_communities, \"create_final_communities.parquet\", storage\n",
-    ")\n",
-    "await write_table_to_storage(\n",
-    "    final_community_reports, \"create_final_community_reports.parquet\", storage\n",
+    "    final_community_reports, \"create_final_community_reports\", storage\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from graphrag.cache.factory import create_cache\n",
+    "from graphrag.cache.factory import CacheFactory\n",
     "from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks\n",
+    "from graphrag.index.config.embeddings import get_embedded_fields, get_embedding_settings\n",
     "from graphrag.index.flows.generate_text_embeddings import generate_text_embeddings\n",
     "\n",
     "# We only need to re-run the embeddings workflow, to ensure that embeddings for all required search fields are in place\n",
     "# We'll construct the context and run this function flow directly to avoid everything else\n",
     "\n",
-    "workflow = next(\n",
-    "    (x for x in pipeline_config.workflows if x.name == \"generate_text_embeddings\"), None\n",
-    ")\n",
-    "config = workflow.config\n",
-    "text_embed = config.get(\"text_embed\", {})\n",
-    "embedded_fields = config.get(\"embedded_fields\", {})\n",
+    "\n",
+    "embedded_fields = get_embedded_fields(config)\n",
+    "text_embed = get_embedding_settings(config)\n",
     "callbacks = NoopWorkflowCallbacks()\n",
-    "cache = create_cache(pipeline_config.cache, PROJECT_DIRECTORY)\n",
+    "cache_config = config.cache.model_dump()  # type: ignore\n",
+    "cache = CacheFactory().create_cache(\n",
+    "    cache_type=cache_config[\"type\"],  # type: ignore\n",
+    "    root_dir=PROJECT_DIRECTORY,\n",
+    "    kwargs=cache_config,\n",
+    ")\n",
     "\n",
     "await generate_text_embeddings(\n",
     "    final_documents=None,\n",

graphrag/api/index.py

@@ -11,7 +11,7 @@
 import logging
 
 from graphrag.cache.noop_pipeline_cache import NoopPipelineCache
-from graphrag.callbacks.factory import create_pipeline_reporter
+from graphrag.callbacks.reporting import create_pipeline_reporter
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
 from graphrag.config.enums import CacheType
 from graphrag.config.models.graph_rag_config import GraphRagConfig
@@ -24,8 +24,6 @@
 
 async def build_index(
     config: GraphRagConfig,
-    run_id: str = "",
-    is_resume_run: bool = False,
     memory_profile: bool = False,
     callbacks: list[WorkflowCallbacks] | None = None,
     progress_logger: ProgressLogger | None = None,
@@ -36,10 +34,6 @@ async def build_index(
     ----------
     config : GraphRagConfig
         The configuration.
-    run_id : str
-        The run id. Creates a output directory with this name.
-    is_resume_run : bool default=False
-        Whether to resume a previous index run.
     memory_profile : bool
         Whether to enable memory profiling.
     callbacks : list[WorkflowCallbacks] | None default=None
@@ -52,11 +46,7 @@ async def build_index(
     list[PipelineRunResult]
         The list of pipeline run results
     """
-    is_update_run = bool(config.update_index_storage)
-
-    if is_resume_run and is_update_run:
-        msg = "Cannot resume and update a run at the same time."
-        raise ValueError(msg)
+    is_update_run = bool(config.update_index_output)
 
     pipeline_cache = (
         NoopPipelineCache() if config.cache.type == CacheType.none is None else None
@@ -78,7 +68,6 @@ async def build_index(
         cache=pipeline_cache,
         callbacks=callbacks,
         logger=progress_logger,
-        run_id=run_id,
         is_update_run=is_update_run,
     ):
         outputs.append(output)

graphrag/api/prompt_tune.py

@@ -17,7 +17,7 @@
 from graphrag.config.models.graph_rag_config import GraphRagConfig
 from graphrag.index.llm.load_llm import load_llm
 from graphrag.logger.print_progress import PrintProgressLogger
-from graphrag.prompt_tune.defaults import MAX_TOKEN_COUNT
+from graphrag.prompt_tune.defaults import MAX_TOKEN_COUNT, PROMPT_TUNING_MODEL_ID
 from graphrag.prompt_tune.generator.community_report_rating import (
     generate_community_report_rating,
 )
@@ -95,9 +95,11 @@ async def generate_indexing_prompts(
     )
 
     # Create LLM from config
+    # TODO: Expose way to specify Prompt Tuning model ID through config
+    default_llm_settings = config.get_language_model_config(PROMPT_TUNING_MODEL_ID)
     llm = load_llm(
         "prompt_tuning",
-        config.llm,
+        default_llm_settings,
         cache=None,
         callbacks=NoopWorkflowCallbacks(),
     )
@@ -120,14 +122,17 @@ async def generate_indexing_prompts(
     )
 
     entity_types = None
+    entity_extraction_llm_settings = config.get_language_model_config(
+        config.entity_extraction.model_id
+    )
     if discover_entity_types:
         logger.info("Generating entity types...")
         entity_types = await generate_entity_types(
             llm,
             domain=domain,
             persona=persona,
             docs=doc_list,
-            json_mode=config.llm.model_supports_json or False,
+            json_mode=entity_extraction_llm_settings.model_supports_json or False,
         )
 
     logger.info("Generating entity relationship examples...")
@@ -147,7 +152,7 @@ async def generate_indexing_prompts(
         examples=examples,
         language=language,
         json_mode=False,  # config.llm.model_supports_json should be used, but these prompts are used in non-json mode by the index engine
-        encoding_model=config.encoding_model,
+        encoding_model=entity_extraction_llm_settings.encoding_model,
         max_token_count=max_tokens,
         min_examples_required=min_examples_required,
     )

graphrag/api/query.py

@@ -24,12 +24,13 @@
 import pandas as pd
 from pydantic import validate_call
 
-from graphrag.config.models.graph_rag_config import GraphRagConfig
-from graphrag.index.config.embeddings import (
+from graphrag.config.embeddings import (
     community_full_content_embedding,
+    create_collection_name,
     entity_description_embedding,
     text_unit_text_embedding,
 )
+from graphrag.config.models.graph_rag_config import GraphRagConfig
 from graphrag.logger.print_progress import PrintProgressLogger
 from graphrag.query.factory import (
     get_basic_search_engine,
@@ -47,7 +48,6 @@
     read_indexer_text_units,
 )
 from graphrag.utils.cli import redact
-from graphrag.utils.embeddings import create_collection_name
 from graphrag.vector_stores.base import BaseVectorStore
 from graphrag.vector_stores.factory import VectorStoreFactory
 
@@ -244,7 +244,7 @@ async def local_search(
     ------
     TODO: Document any exceptions to expect.
     """
-    vector_store_args = config.embeddings.vector_store
+    vector_store_args = config.vector_store.model_dump()
     logger.info(f"Vector Store Args: {redact(vector_store_args)}")  # type: ignore # noqa
 
     description_embedding_store = _get_embedding_store(
@@ -310,7 +310,7 @@ async def local_search_streaming(
     ------
     TODO: Document any exceptions to expect.
     """
-    vector_store_args = config.embeddings.vector_store
+    vector_store_args = config.vector_store.model_dump()
     logger.info(f"Vector Store Args: {redact(vector_store_args)}")  # type: ignore # noqa
 
     description_embedding_store = _get_embedding_store(
@@ -381,7 +381,7 @@ async def drift_search_streaming(
     ------
     TODO: Document any exceptions to expect.
     """
-    vector_store_args = config.embeddings.vector_store
+    vector_store_args = config.vector_store.model_dump()
     logger.info(f"Vector Store Args: {redact(vector_store_args)}")  # type: ignore # noqa
 
     description_embedding_store = _get_embedding_store(
@@ -465,7 +465,7 @@ async def drift_search(
     ------
     TODO: Document any exceptions to expect.
     """
-    vector_store_args = config.embeddings.vector_store
+    vector_store_args = config.vector_store.model_dump()
     logger.info(f"Vector Store Args: {redact(vector_store_args)}")  # type: ignore # noqa
 
     description_embedding_store = _get_embedding_store(
@@ -531,7 +531,7 @@ async def basic_search(
     ------
     TODO: Document any exceptions to expect.
     """
-    vector_store_args = config.embeddings.vector_store
+    vector_store_args = config.vector_store.model_dump()
     logger.info(f"Vector Store Args: {redact(vector_store_args)}")  # type: ignore # noqa
 
     description_embedding_store = _get_embedding_store(
@@ -576,7 +576,7 @@ async def basic_search_streaming(
     ------
     TODO: Document any exceptions to expect.
     """
-    vector_store_args = config.embeddings.vector_store
+    vector_store_args = config.vector_store.model_dump()
     logger.info(f"Vector Store Args: {redact(vector_store_args)}")  # type: ignore # noqa
 
     description_embedding_store = _get_embedding_store(