Add user-provided entity types to prompt tuning

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

@@ -0,0 +1,4 @@
+{
+  "type": "minor",
+  "description": "Add support for user-provided entity types during prompt tuning."
+}

packages/graphrag/graphrag/api/prompt_tune.py

@@ -48,6 +48,40 @@
 logger = logging.getLogger(__name__)
 
 
+def _normalize_entity_types(entity_types: list[str] | None) -> list[str] | None:
+    """Normalize user-provided entity types."""
+    if not entity_types:
+        return None
+
+    normalized = []
+    for entity_type in entity_types:
+        value = str(entity_type).strip()
+        if value and value not in normalized:
+            normalized.append(value)
+
+    return normalized or None
+
+
+def _merge_entity_types(
+    provided_entity_types: list[str] | None,
+    generated_entity_types: list[str] | str | None,
+) -> list[str] | None:
+    """Merge user-provided and generated entity types, preserving order."""
+    if isinstance(generated_entity_types, str):
+        generated_entity_types = [generated_entity_types]
+
+    merged = []
+    for entity_type in [
+        *(provided_entity_types or []),
+        *(generated_entity_types or []),
+    ]:
+        value = str(entity_type).strip()
+        if value and value not in merged:
+            merged.append(value)
+
+    return merged or None
+
+
 @validate_call(config={"arbitrary_types_allowed": True})
 async def generate_indexing_prompts(
     config: GraphRagConfig,
@@ -57,6 +91,7 @@ async def generate_indexing_prompts(
     language: str | None = None,
     max_tokens: int = MAX_TOKEN_COUNT,
     discover_entity_types: bool = True,
+    entity_types: list[str] | None = None,
     min_examples_required: PositiveInt = 2,
     n_subset_max: PositiveInt = 300,
     k: PositiveInt = 15,
@@ -75,6 +110,7 @@ async def generate_indexing_prompts(
     - language: The language to use for the prompts.
     - max_tokens: The maximum number of tokens to use on entity extraction prompts
     - discover_entity_types: Generate entity types.
+    - entity_types: Optional user-provided entity types to include in the generated prompts.
     - min_examples_required: The minimum number of examples required for entity extraction prompts.
     - n_subset_max: The number of text chunks to embed when using auto selection method.
     - k: The number of documents to select when using auto selection method.
@@ -120,20 +156,26 @@ async def generate_indexing_prompts(
         llm, domain=domain, persona=persona, docs=doc_list
     )
 
-    entity_types = None
+    provided_entity_types = _normalize_entity_types(entity_types)
+    generated_entity_types = None
     extract_graph_llm_settings = config.get_completion_model_config(
         config.extract_graph.completion_model_id
     )
     if discover_entity_types:
         logger.info("Generating entity types...")
-        entity_types = await generate_entity_types(
+        generated_entity_types = await generate_entity_types(
             llm,
             domain=domain,
             persona=persona,
             docs=doc_list,
             json_mode=True,
         )
 
+    entity_types = _merge_entity_types(
+        provided_entity_types=provided_entity_types,
+        generated_entity_types=generated_entity_types,
+    )
+
     logger.info("Generating entity relationship examples...")
     examples = await generate_entity_relationship_examples(
         llm,

packages/graphrag/graphrag/cli/main.py

@@ -325,6 +325,14 @@ def _prompt_tune_cli(
         "--discover-entity-types/--no-discover-entity-types",
         help="Discover and extract unspecified entity types.",
     ),
+    entity_types: str | None = typer.Option(
+        None,
+        "--entity-types",
+        help=(
+            "Comma-separated list of entity types to include in generated prompts. "
+            "When discovery is enabled, these are merged with discovered entity types."
+        ),
+    ),
     output: Path = typer.Option(
         Path("prompts"),
         "--output",
@@ -353,6 +361,7 @@ def _prompt_tune_cli(
             overlap=overlap,
             language=language,
             discover_entity_types=discover_entity_types,
+            entity_types=entity_types,
             output=output,
             n_subset_max=n_subset_max,
             k=k,

packages/graphrag/graphrag/cli/prompt_tune.py

@@ -22,6 +22,17 @@
 logger = logging.getLogger(__name__)
 
 
+def _parse_entity_types(entity_types: str | None) -> list[str] | None:
+    """Parse a comma-separated list of entity types."""
+    if not entity_types:
+        return None
+
+    parsed = [entity_type.strip() for entity_type in entity_types.split(",")]
+    parsed = [entity_type for entity_type in parsed if entity_type]
+
+    return parsed or None
+
+
 async def prompt_tune(
     root: Path,
     domain: str | None,
@@ -33,6 +44,7 @@ async def prompt_tune(
     overlap: int,
     language: str | None,
     discover_entity_types: bool,
+    entity_types: str | None,
     output: Path,
     n_subset_max: int,
     k: int,
@@ -51,6 +63,7 @@ async def prompt_tune(
     - chunk_size: The chunk token size to use.
     - language: The language to use for the prompts.
     - discover_entity_types: Generate entity types.
+    - entity_types: Comma-separated list of entity types to include in generated prompts.
     - output: The output folder to store the prompts.
     - n_subset_max: The number of text chunks to embed when using auto selection method.
     - k: The number of documents to select when using auto selection method.
@@ -87,6 +100,7 @@ async def prompt_tune(
         language=language,
         max_tokens=max_tokens,
         discover_entity_types=discover_entity_types,
+        entity_types=_parse_entity_types(entity_types),
         min_examples_required=min_examples_required,
         n_subset_max=n_subset_max,
         k=k,

tests/unit/graphrag/prompt_tune/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License

tests/unit/graphrag/prompt_tune/test_entity_types.py

@@ -0,0 +1,50 @@
+# Copyright (c) 2024 Microsoft Corporation.
+# Licensed under the MIT License
+
+from graphrag.api.prompt_tune import _merge_entity_types, _normalize_entity_types
+from graphrag.cli.prompt_tune import _parse_entity_types
+
+
+def test_parse_entity_types_splits_commas_and_ignores_blanks():
+    assert _parse_entity_types("person, organization,, event ") == [
+        "person",
+        "organization",
+        "event",
+    ]
+
+
+def test_parse_entity_types_returns_none_for_empty_input():
+    assert _parse_entity_types(None) is None
+    assert _parse_entity_types("") is None
+    assert _parse_entity_types(" , ") is None
+
+
+def test_normalize_entity_types_strips_blanks_and_deduplicates():
+    assert _normalize_entity_types([" person ", "", "organization", "person"]) == [
+        "person",
+        "organization",
+    ]
+
+
+def test_merge_entity_types_preserves_user_order_and_deduplicates_generated_values():
+    assert _merge_entity_types(
+        provided_entity_types=["person", "organization"],
+        generated_entity_types=["organization", "event"],
+    ) == ["person", "organization", "event"]
+
+
+def test_merge_entity_types_accepts_string_generated_entity_types():
+    assert _merge_entity_types(
+        provided_entity_types=["person"],
+        generated_entity_types="event",
+    ) == ["person", "event"]
+
+
+def test_merge_entity_types_returns_none_when_no_entity_types_are_available():
+    assert (
+        _merge_entity_types(
+            provided_entity_types=None,
+            generated_entity_types=None,
+        )
+        is None
+    )