Improve CLI speed with lazy imports (#1319)

Author: jgbradley1

.semversioner/next-release/patch-20241025031711368197.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "move import statements out of init files"
+}

.semversioner/next-release/patch-20241031180003172666.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "fix autocompletion of existing files/directory paths."
+}

docs/prompt_tuning/auto_prompt_tuning.md

@@ -20,9 +20,9 @@ Before running auto tuning, ensure you have already initialized your workspace w
 You can run the main script from the command line with various options:
 
 ```bash
-graphrag prompt-tune [--root ROOT] [--domain DOMAIN]  [--method METHOD] [--limit LIMIT] [--language LANGUAGE] \
+graphrag prompt-tune [--root ROOT] [--config CONFIG] [--domain DOMAIN]  [--selection-method METHOD] [--limit LIMIT] [--language LANGUAGE] \
 [--max-tokens MAX_TOKENS] [--chunk-size CHUNK_SIZE] [--n-subset-max N_SUBSET_MAX] [--k K] \
-[--min-examples-required MIN_EXAMPLES_REQUIRED] [--no-entity-types] [--output OUTPUT]
+[--min-examples-required MIN_EXAMPLES_REQUIRED] [--discover-entity-types] [--output OUTPUT]
 ```
 
 ## Command-Line Options
@@ -49,7 +49,7 @@ graphrag prompt-tune [--root ROOT] [--domain DOMAIN]  [--method METHOD] [--limit
 
 - `--min-examples-required` (optional): The minimum number of examples required for entity extraction prompts. Default is 2.
 
-- `--no-entity-types` (optional): Use untyped entity extraction generation. We recommend using this when your data covers a lot of topics or it is highly randomized.
+- `--discover-entity-types` (optional): Allow the LLM to discover and extract entities automatically. We recommend using this when your data covers a lot of topics or it is highly randomized.
 
 - `--output` (optional): The folder to save the generated prompts. Default is "prompts".
 

examples/custom_input/run.py

@@ -5,7 +5,7 @@
 
 import pandas as pd
 
-from graphrag.index import run_pipeline_with_config
+from graphrag.index.run import run_pipeline_with_config
 
 pipeline_file = os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "./pipeline.yml"

examples/single_verb/run.py

@@ -5,8 +5,8 @@
 
 import pandas as pd
 
-from graphrag.index import run_pipeline, run_pipeline_with_config
-from graphrag.index.config import PipelineWorkflowReference
+from graphrag.index.config.workflow import PipelineWorkflowReference
+from graphrag.index.run import run_pipeline, run_pipeline_with_config
 
 # our fake dataset
 dataset = pd.DataFrame([{"col1": 2, "col2": 4}, {"col1": 5, "col2": 10}])

examples/use_built_in_workflows/run.py

@@ -3,9 +3,10 @@
 import asyncio
 import os
 
-from graphrag.index import run_pipeline, run_pipeline_with_config
-from graphrag.index.config import PipelineCSVInputConfig, PipelineWorkflowReference
-from graphrag.index.input import load_input
+from graphrag.index.config.input import PipelineCSVInputConfig
+from graphrag.index.config.workflow import PipelineWorkflowReference
+from graphrag.index.input.load_input import load_input
+from graphrag.index.run import run_pipeline, run_pipeline_with_config
 
 sample_data_dir = os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "../_sample_data/"

graphrag/__main__.py

@@ -3,6 +3,6 @@
 
 """The GraphRAG package."""
 
-from .cli.main import app
+from graphrag.cli.main import app
 
 app(prog_name="graphrag")

graphrag/api/__init__.py

@@ -8,14 +8,15 @@
 """
 
 from graphrag.api.index import build_index
-from graphrag.api.prompt_tune import DocSelectionType, generate_indexing_prompts
+from graphrag.api.prompt_tune import generate_indexing_prompts
 from graphrag.api.query import (
     drift_search,
     global_search,
     global_search_streaming,
     local_search,
     local_search_streaming,
 )
+from graphrag.prompt_tune.types import DocSelectionType
 
 __all__ = [  # noqa: RUF022
     # index API

graphrag/api/index.py

@@ -10,13 +10,14 @@
 
 from pathlib import Path
 
-from graphrag.config import CacheType, GraphRagConfig
+from graphrag.config.enums import CacheType
+from graphrag.config.models.graph_rag_config import GraphRagConfig
 from graphrag.index.cache.noop_pipeline_cache import NoopPipelineCache
 from graphrag.index.create_pipeline_config import create_pipeline_config
 from graphrag.index.emit.types import TableEmitterType
 from graphrag.index.run import run_pipeline_with_config
 from graphrag.index.typing import PipelineRunResult
-from graphrag.logging import ProgressReporter
+from graphrag.logging.base import ProgressReporter
 from graphrag.vector_stores.factory import VectorStoreType
 
 

graphrag/api/prompt_tune.py

@@ -15,25 +15,32 @@
 from pydantic import PositiveInt, validate_call
 
 from graphrag.config.models.graph_rag_config import GraphRagConfig
-from graphrag.index.llm import load_llm
-from graphrag.logging import PrintProgressReporter
-from graphrag.prompt_tune.generator import (
-    MAX_TOKEN_COUNT,
-    create_community_summarization_prompt,
-    create_entity_extraction_prompt,
-    create_entity_summarization_prompt,
-    detect_language,
+from graphrag.index.llm.load_llm import load_llm
+from graphrag.logging.print_progress import PrintProgressReporter
+from graphrag.prompt_tune.defaults import MAX_TOKEN_COUNT
+from graphrag.prompt_tune.generator.community_report_rating import (
     generate_community_report_rating,
+)
+from graphrag.prompt_tune.generator.community_report_summarization import (
+    create_community_summarization_prompt,
+)
+from graphrag.prompt_tune.generator.community_reporter_role import (
     generate_community_reporter_role,
-    generate_domain,
+)
+from graphrag.prompt_tune.generator.domain import generate_domain
+from graphrag.prompt_tune.generator.entity_extraction_prompt import (
+    create_entity_extraction_prompt,
+)
+from graphrag.prompt_tune.generator.entity_relationship import (
     generate_entity_relationship_examples,
-    generate_entity_types,
-    generate_persona,
 )
-from graphrag.prompt_tune.loader import (
-    MIN_CHUNK_SIZE,
-    load_docs_in_chunks,
+from graphrag.prompt_tune.generator.entity_summarization_prompt import (
+    create_entity_summarization_prompt,
 )
+from graphrag.prompt_tune.generator.entity_types import generate_entity_types
+from graphrag.prompt_tune.generator.language import detect_language
+from graphrag.prompt_tune.generator.persona import generate_persona
+from graphrag.prompt_tune.loader.input import MIN_CHUNK_SIZE, load_docs_in_chunks
 from graphrag.prompt_tune.types import DocSelectionType