diff --git a/.semversioner/next-release/minor-20260705140348701022.json b/.semversioner/next-release/minor-20260705140348701022.json new file mode 100644 index 0000000000..3dc06911d0 --- /dev/null +++ b/.semversioner/next-release/minor-20260705140348701022.json @@ -0,0 +1,4 @@ +{ + "type": "minor", + "description": "Add support for user-provided entity types during prompt tuning." +} diff --git a/packages/graphrag/graphrag/api/prompt_tune.py b/packages/graphrag/graphrag/api/prompt_tune.py index e1741f014b..b4eba96e66 100644 --- a/packages/graphrag/graphrag/api/prompt_tune.py +++ b/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,13 +156,14 @@ 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, @@ -134,6 +171,11 @@ async def generate_indexing_prompts( 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, diff --git a/packages/graphrag/graphrag/cli/main.py b/packages/graphrag/graphrag/cli/main.py index 5259cf4b85..eb60b110b2 100644 --- a/packages/graphrag/graphrag/cli/main.py +++ b/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, diff --git a/packages/graphrag/graphrag/cli/prompt_tune.py b/packages/graphrag/graphrag/cli/prompt_tune.py index 249dd7cd07..124fc96fc4 100644 --- a/packages/graphrag/graphrag/cli/prompt_tune.py +++ b/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, diff --git a/tests/unit/graphrag/prompt_tune/__init__.py b/tests/unit/graphrag/prompt_tune/__init__.py new file mode 100644 index 0000000000..0a3e38adfb --- /dev/null +++ b/tests/unit/graphrag/prompt_tune/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) 2024 Microsoft Corporation. +# Licensed under the MIT License diff --git a/tests/unit/graphrag/prompt_tune/test_entity_types.py b/tests/unit/graphrag/prompt_tune/test_entity_types.py new file mode 100644 index 0000000000..703755fc0d --- /dev/null +++ b/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 + )