diff --git a/docs-site/src/content/docs/lexical-graph/composable-extraction-pipeline-usage.md b/docs-site/src/content/docs/lexical-graph/composable-extraction-pipeline-usage.md new file mode 100644 index 000000000..215618ed6 --- /dev/null +++ b/docs-site/src/content/docs/lexical-graph/composable-extraction-pipeline-usage.md @@ -0,0 +1,246 @@ +# Composable Extraction Pipeline — Usage Guide + +The composable extraction pipeline lets you customize how the lexical graph extracts entities, relationships, and topics from your documents. Instead of the fixed default pipeline, you define a sequence of **stages** that are validated at build time and composed into a processing pipeline. + +## Quick Start + +```python +from graphrag_toolkit.lexical_graph.indexing.extract import ( + LLMPropositionStage, + LLMTopicExtractionStage, + PipelineBuilder, +) + +# Build a pipeline identical to the default +builder = PipelineBuilder() +builder.add(LLMPropositionStage()) +builder.add(LLMTopicExtractionStage()) +transforms = builder.build() +``` + +Or use `ExtractionConfig.from_stages()` for a custom composable pipeline: + +```python +from graphrag_toolkit.lexical_graph.lexical_graph_index import ExtractionConfig, LexicalGraphIndex + +config = ExtractionConfig.from_stages( + stages=[ + LLMPropositionStage(), + LLMTopicExtractionStage(), + ] +) +``` + +When `stages` is provided, the pipeline builder validates that each stage's required inputs are satisfied by prior stages' outputs. If not, it raises a `ValueError` at build time — no silent failures at runtime. + +## Available Stages + +### Proposition Extraction + +| Stage | Description | Requires | +|-------|-------------|----------| +| `LLMPropositionStage()` | Extracts propositions using an LLM | — | +| `LocalPropositionStage()` | Extracts propositions using a local transformer model | — | +| `BatchLLMPropositionStage(batch_config)` | Batch proposition extraction via Bedrock | `BatchConfig` | + +### Topic & Entity Extraction + +| Stage | Description | Requires | +|-------|-------------|----------| +| `LLMTopicExtractionStage()` | Extracts topics, entities, and relationships from propositions | propositions | +| `LLMTopicExtractionStage(use_propositions=False)` | Extracts directly from source text | — | +| `BatchTopicExtractionStage(batch_config)` | Batch topic extraction via Bedrock | `BatchConfig` | + +### NER (Named Entity Recognition) + +| Stage | Description | Requires | +|-------|-------------|----------| +| `NERExtractionStage(entity_labels=[...])` | CPU-based NER using GLiNER | `pip install gliner` | +| `EntityMergeStage()` | Merges NER entities into LLM-extracted topics | NER output + topics | + +### Filtering + +| Stage | Description | Requires | +|-------|-------------|----------| +| `SchemaFilterStage(schema)` | Filters entities/relationships against a schema | topics | + +## Defining a Schema + +Use `ExtractionSchema` to constrain what entity types and relationship types are extracted: + +```python +from graphrag_toolkit.lexical_graph.indexing.extract import ( + ExtractionSchema, + EntityTypeConfig, + SchemaFilterStage, +) + +schema = ExtractionSchema( + entity_types={ + 'Person': EntityTypeConfig( + description='A human individual', + attributes=['age', 'role'], + aliases=['Individual', 'Human'], + ), + 'Company': EntityTypeConfig( + description='A business organization', + ), + }, + relationship_types=['WORKS_FOR', 'FOUNDED', 'ACQUIRED'], + strict=True, # When True, SchemaFilterStage removes non-matching entities/relationships +) +``` + +Add the filter after topic extraction: + +```python +config = ExtractionConfig.from_stages( + stages=[ + LLMPropositionStage(), + LLMTopicExtractionStage(), + SchemaFilterStage(schema), + ] +) +``` + +When `strict=False` (the default), the filter stage passes everything through unchanged — useful for soft constraints where you want the schema for prompt guidance but not hard filtering. + +When `schema` is provided to `ExtractionConfig.from_stages()`, it is automatically injected into the `LLMTopicExtractionStage` prompt — guiding the LLM to focus on your entity types and relationships during extraction, not just during filtering. + +## Combining NER with LLM Extraction + +Augment LLM extraction with GLiNER-based NER to catch entities the LLM might miss: + +```python +from graphrag_toolkit.lexical_graph.indexing.extract import ( + LLMPropositionStage, + LLMTopicExtractionStage, + NERExtractionStage, + EntityMergeStage, +) + +config = ExtractionConfig.from_stages( + stages=[ + LLMPropositionStage(), + NERExtractionStage( + entity_labels=['Person', 'Organization', 'Location'], + threshold=0.5, + ), + LLMTopicExtractionStage(), + EntityMergeStage(fuzzy_threshold=0.85), # Merges NER entities, fuzzy dedup + ] +) +``` + +`EntityMergeStage` performs case-insensitive deduplication — if the LLM already extracted "John" and NER also found "john", it won't be added twice. + +For fuzzy deduplication (e.g., "DataBridge" ≈ "DataBridge AI"), set a similarity threshold: + +```python +EntityMergeStage(fuzzy_threshold=0.85) # 0-1, higher = stricter matching +``` + +With `fuzzy_threshold=None` (default), only exact case-insensitive matches are deduplicated. + +## Using JSON Output for Topic Extraction + +The `TopicExtractor` now supports JSON-formatted LLM output, which is more reliable than text parsing: + +```python +from graphrag_toolkit.lexical_graph.indexing.extract import TopicExtractor + +extractor = TopicExtractor(output_format='json') +``` + +When `output_format='json'`, the extractor tries JSON parsing first and falls back to the legacy text parser if JSON parsing fails. + +## Using Local Proposition Extraction + +Skip the LLM for proposition extraction and use a local transformer model instead: + +```python +from graphrag_toolkit.lexical_graph.indexing.extract import ( + LocalPropositionStage, + LLMTopicExtractionStage, +) + +config = ExtractionConfig.from_stages( + stages=[ + LocalPropositionStage(), # Runs locally, no LLM cost + LLMTopicExtractionStage(), + ] +) +``` + +## Full Example: Schema-Constrained Pipeline with NER + +```python +from graphrag_toolkit.lexical_graph.lexical_graph_index import ExtractionConfig +from graphrag_toolkit.lexical_graph.indexing.extract import ( + LLMPropositionStage, + LLMTopicExtractionStage, + NERExtractionStage, + EntityMergeStage, + SchemaFilterStage, + ExtractionSchema, + EntityTypeConfig, +) + +schema = ExtractionSchema( + entity_types={ + 'Person': EntityTypeConfig(aliases=['Individual']), + 'Company': EntityTypeConfig(aliases=['Organization', 'Corp']), + 'Technology': EntityTypeConfig(), + }, + relationship_types=['WORKS_FOR', 'USES', 'DEVELOPS', 'ACQUIRED'], + strict=True, +) + +config = ExtractionConfig.from_stages( + stages=[ + LLMPropositionStage(), + NERExtractionStage(entity_labels=['Person', 'Organization', 'Technology']), + LLMTopicExtractionStage(), + EntityMergeStage(fuzzy_threshold=0.85), + SchemaFilterStage(schema), + ], + schema=schema, +) +``` + +## Writing Custom Stages + +Implement the `ExtractionStage` ABC: + +```python +from typing import List +from llama_index.core.schema import BaseNode, TransformComponent +from graphrag_toolkit.lexical_graph.indexing.extract import ExtractionStage + +class MyCustomTransform(TransformComponent): + def __call__(self, nodes, **kwargs): + for node in nodes: + # Your processing logic here + node.metadata['my_key'] = 'my_value' + return nodes + +class MyCustomStage(ExtractionStage): + def input_keys(self) -> List[str]: + return [] # Keys this stage requires in node metadata + + def output_keys(self) -> List[str]: + return ['my_key'] # Keys this stage adds to node metadata + + def as_transform(self) -> TransformComponent: + return MyCustomTransform() + + @property + def stage_type(self) -> str: + return 'custom' +``` + +The `PipelineBuilder` validates that `input_keys()` are available (from initial context or prior stages' `output_keys()`) before the pipeline runs. This catches wiring errors at build time. + +## Backward Compatibility + +When `ExtractionConfig.from_stages()` is not used, the existing default pipeline (LLM propositions → topic extraction) is used unchanged. All existing code continues to work without modification. diff --git a/examples/lexical-graph-local-dev/README.md b/examples/lexical-graph-local-dev/README.md index bb1f2f925..b47c7cd86 100644 --- a/examples/lexical-graph-local-dev/README.md +++ b/examples/lexical-graph-local-dev/README.md @@ -14,6 +14,7 @@ This example provides a complete local development environment for the GraphRAG - [**03-Querying-with-Prompting**](./notebooks/03-Querying-with-Prompting.ipynb) – Advanced querying with custom prompts and prompt providers - [**04-Advanced-Configuration-Examples**](./notebooks/04-Advanced-Configuration-Examples.ipynb) – Advanced reader configurations and metadata handling - [**05-S3-Directory-Reader-Provider**](./notebooks/05-S3-Directory-Reader-Provider.ipynb) – S3-based document reading and processing +- [**06-Composable-Extraction-Pipeline**](./notebooks/06-Composable-Extraction-Pipeline.ipynb) – Customizable extraction stages, schema filtering, and NER integration ## Quick Start diff --git a/examples/lexical-graph-local-dev/notebooks/06-Composable-Extraction-Pipeline.ipynb b/examples/lexical-graph-local-dev/notebooks/06-Composable-Extraction-Pipeline.ipynb new file mode 100644 index 000000000..482d372c3 --- /dev/null +++ b/examples/lexical-graph-local-dev/notebooks/06-Composable-Extraction-Pipeline.ipynb @@ -0,0 +1,405 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 06 - Composable Extraction Pipeline\n", + "\n", + "This notebook demonstrates how to customize the extraction pipeline \u2014 the process that turns your documents into entities, relationships, and topics for the knowledge graph.\n", + "\n", + "By default, lexical-graph runs a fixed two-step extraction (propositions \u2192 topics). The composable pipeline lets you define your own sequence of stages, add schema constraints, combine LLM extraction with local NER models, and write custom processing steps." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup\n", + "\n", + "If you haven't already, install the toolkit and dependencies using the [Setup](./00-Setup.ipynb) notebook." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Install optional dependency for NER stage (Example 4)\n", + "# Skip this cell if you don't plan to use NER\n", + "%pip install -q gliner" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%reload_ext dotenv\n", + "%dotenv\n", + "\n", + "import os\n", + "\n", + "from graphrag_toolkit.lexical_graph import LexicalGraphIndex, set_logging_config\n", + "from graphrag_toolkit.lexical_graph.lexical_graph_index import ExtractionConfig\n", + "from graphrag_toolkit.lexical_graph.indexing.extract import (\n", + " PipelineBuilder,\n", + " ExtractionSchema,\n", + " EntityTypeConfig,\n", + " LLMPropositionStage,\n", + " LocalPropositionStage,\n", + " LLMTopicExtractionStage,\n", + " SchemaFilterStage,\n", + " NERExtractionStage,\n", + " EntityMergeStage,\n", + ")\n", + "\n", + "set_logging_config('INFO')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example 1: Pipeline Variations\n", + "\n", + "The simplest use \u2014 pick which stages run and in what order. Three common patterns:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Pattern A: Default \u2014 same as built-in behavior, but explicit\n", + "default_pipeline = ExtractionConfig.from_stages(\n", + " stages=[\n", + " LLMPropositionStage(),\n", + " LLMTopicExtractionStage(),\n", + " ]\n", + ")\n", + "\n", + "# Pattern B: Skip propositions \u2014 extract topics directly from text (faster, cheaper)\n", + "direct_pipeline = ExtractionConfig.from_stages(\n", + " stages=[\n", + " LLMTopicExtractionStage(use_propositions=False),\n", + " ]\n", + ")\n", + "\n", + "# Pattern C: Local propositions \u2014 use a local model instead of LLM (free, no API calls)\n", + "local_pipeline = ExtractionConfig.from_stages(\n", + " stages=[\n", + " LocalPropositionStage(),\n", + " LLMTopicExtractionStage(),\n", + " ]\n", + ")\n", + "\n", + "for name, config in [('Default', default_pipeline), ('Direct', direct_pipeline), ('Local', local_pipeline)]:\n", + " stages = [type(s).__name__ for s in config.stages]\n", + " print(f'{name:8s} \u2192 {\" \u2192 \".join(stages)}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example 2: Schema-Constrained Extraction\n", + "\n", + "Define which entity types and relationship types you care about. When `schema` is passed to\n", + "`ExtractionConfig`, it is automatically injected into the LLM prompt \u2014 guiding extraction, not\n", + "just filtering. The `SchemaFilterStage` then removes anything the LLM still extracted outside the schema." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "schema = ExtractionSchema(\n", + " entity_types={\n", + " 'Person': EntityTypeConfig(\n", + " description='A human individual',\n", + " aliases=['Individual', 'Human'], # These also match\n", + " ),\n", + " 'Company': EntityTypeConfig(\n", + " description='A business organization',\n", + " aliases=['Organization', 'Corp'],\n", + " ),\n", + " },\n", + " relationship_types=['WORKS_FOR', 'FOUNDED', 'ACQUIRED'],\n", + " strict=True, # Enforce filtering\n", + ")\n", + "\n", + "# See what the schema looks like as a prompt constraint\n", + "print('Prompt constraint that gets injected into the LLM:\\n')\n", + "print(schema.format_as_prompt_constraint())\n", + "\n", + "# Use it in a pipeline\n", + "extraction_config = ExtractionConfig.from_stages(\n", + " stages=[\n", + " LLMPropositionStage(),\n", + " LLMTopicExtractionStage(),\n", + " SchemaFilterStage(schema),\n", + " ],\n", + " schema=schema,\n", + ")\n", + "\n", + "print(f'\\nAllowed entity types: {schema.entity_type_names()}')\n", + "print(f'Allowed relationships: {schema.relationship_types}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example 3: Pipeline Validation\n", + "\n", + "The `PipelineBuilder` checks that each stage's required inputs are satisfied by prior stages' outputs. This catches wiring mistakes before you spend time and money on LLM calls." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Valid pipeline \u2014 propositions come before topic extraction\n", + "builder = PipelineBuilder()\n", + "builder.add(LLMPropositionStage())\n", + "builder.add(LLMTopicExtractionStage())\n", + "transforms = builder.build()\n", + "print(f'\u2705 Valid pipeline: {len(transforms)} transforms')\n", + "print(f' Available keys after build: {builder.available_keys}')\n", + "\n", + "# Invalid pipeline \u2014 topic extraction needs propositions but none provided\n", + "print()\n", + "try:\n", + " PipelineBuilder().add(LLMTopicExtractionStage()).build()\n", + "except ValueError as e:\n", + " print(f'\u274c Caught wiring error: {e}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example 4: NER + LLM Hybrid Extraction\n", + "\n", + "Augment LLM extraction with a local NER model ([GLiNER](https://github.com/urchade/GLiNER)). NER runs on CPU, costs nothing, and catches entities the LLM might miss. The `EntityMergeStage` combines both sources and deduplicates." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "schema = ExtractionSchema(\n", + " entity_types={\n", + " 'Person': EntityTypeConfig(aliases=['Individual']),\n", + " 'Company': EntityTypeConfig(aliases=['Organization', 'Corp']),\n", + " 'Technology': EntityTypeConfig(),\n", + " },\n", + " relationship_types=['WORKS_FOR', 'USES', 'DEVELOPS', 'ACQUIRED'],\n", + " strict=True,\n", + ")\n", + "\n", + "extraction_config = ExtractionConfig.from_stages(\n", + " stages=[\n", + " LLMPropositionStage(), # 1. Break text into propositions\n", + " NERExtractionStage( # 2. Find entities locally (free)\n", + " entity_labels=['Person', 'Organization', 'Technology'],\n", + " threshold=0.5,\n", + " ),\n", + " LLMTopicExtractionStage(), # 3. Extract topics via LLM\n", + " EntityMergeStage(fuzzy_threshold=0.85), # 4. Merge NER + LLM entities (fuzzy dedup)\n", + " SchemaFilterStage(schema), # 5. Keep only what we want\n", + " ],\n", + " schema=schema,\n", + ")\n", + "\n", + "print('Pipeline:')\n", + "for i, stage in enumerate(extraction_config.stages):\n", + " print(f' {i+1}. {type(stage).__name__:30s} type={stage.stage_type:10s} in={stage.input_keys()} out={stage.output_keys()}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example 5: Custom Stage\n", + "\n", + "Implement the `ExtractionStage` ABC to add your own processing logic \u2014 sentiment analysis, language detection, external NER injection, or anything else." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from typing import List, Sequence\n", + "from llama_index.core.schema import BaseNode, TransformComponent\n", + "from graphrag_toolkit.lexical_graph.indexing.extract import ExtractionStage\n", + "\n", + "STATS_KEY = 'my::text_stats'\n", + "\n", + "class TextStatsTransform(TransformComponent):\n", + " \"\"\"Adds word/char counts to node metadata.\"\"\"\n", + " def __call__(self, nodes: Sequence[BaseNode], **kwargs) -> Sequence[BaseNode]:\n", + " for node in nodes:\n", + " text = node.get_content()\n", + " node.metadata[STATS_KEY] = {'words': len(text.split()), 'chars': len(text)}\n", + " return nodes\n", + "\n", + "class TextStatsStage(ExtractionStage):\n", + " def input_keys(self) -> List[str]:\n", + " return []\n", + " def output_keys(self) -> List[str]:\n", + " return [STATS_KEY]\n", + " def as_transform(self) -> TransformComponent:\n", + " return TextStatsTransform()\n", + " @property\n", + " def stage_type(self) -> str:\n", + " return 'custom'\n", + "\n", + "# Plug it into a pipeline\n", + "builder = PipelineBuilder()\n", + "builder.add(TextStatsStage())\n", + "builder.add(LLMPropositionStage())\n", + "builder.add(LLMTopicExtractionStage())\n", + "transforms = builder.build()\n", + "\n", + "print(f'Pipeline with custom stage: {len(transforms)} transforms')\n", + "print(f'Available keys: {builder.available_keys}')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Example 6: End-to-End \u2014 NER + Schema Filtering\n", + "\n", + "Put it all together with a real file. We load `artifacts/sample.md` (a markdown formatting guide\n", + "mentioning tools like Pandoc and Python). NER and schema are aligned on the same entity types,\n", + "with the schema doing light filtering \u2014 keeping the three main types while dropping anything\n", + "the LLM might hallucinate outside those categories." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from llama_index.core import SimpleDirectoryReader\n", + "\n", + "# Schema aligned with NER labels \u2014 both target the same entity types\n", + "# Light filtering: keeps Software, Format, and Concept; drops anything else\n", + "schema = ExtractionSchema(\n", + " entity_types={\n", + " 'Software': EntityTypeConfig(\n", + " description='A software tool, language, or library',\n", + " aliases=['Tool', 'Language', 'Library', 'Program'],\n", + " ),\n", + " 'Format': EntityTypeConfig(\n", + " description='A file format, markup language, or data notation',\n", + " aliases=['Markup', 'Syntax', 'Notation'],\n", + " ),\n", + " 'Concept': EntityTypeConfig(\n", + " description='A technical concept or standard',\n", + " aliases=['Standard', 'Feature', 'Specification'],\n", + " ),\n", + " },\n", + " relationship_types=['SUPPORTS', 'CONVERTS_TO', 'USES', 'PART_OF', 'RELATED_TO'],\n", + " strict=True,\n", + ")\n", + "\n", + "extraction_config = ExtractionConfig.from_stages(\n", + " stages=[\n", + " LLMPropositionStage(),\n", + " NERExtractionStage(\n", + " # Same types as schema \u2014 NER and schema are aligned\n", + " entity_labels=['Software', 'Format', 'Concept'],\n", + " threshold=0.3,\n", + " ),\n", + " LLMTopicExtractionStage(),\n", + " EntityMergeStage(fuzzy_threshold=0.85), # Fuzzy dedup: 'DataBridge' \u2248 'DataBridge AI'\n", + " SchemaFilterStage(schema), # Light filter: drops only out-of-schema types\n", + " ],\n", + " schema=schema,\n", + ")\n", + "\n", + "graph_index = LexicalGraphIndex(\n", + " os.environ['GRAPH_STORE'],\n", + " os.environ['VECTOR_STORE'],\n", + " indexing_config=extraction_config,\n", + ")\n", + "\n", + "print('Pipeline stages:')\n", + "for i, comp in enumerate(graph_index.extraction_components):\n", + " print(f' {i+1}. {type(comp).__name__}')\n", + "\n", + "# Load sample.md\n", + "docs = SimpleDirectoryReader(input_files=['artifacts/sample.md']).load_data()\n", + "print(f'\\nLoaded {len(docs)} document(s) from artifacts/sample.md')\n", + "\n", + "# Run extraction and build\n", + "graph_index.extract_and_build(docs, show_progress=True)\n", + "\n", + "# Show graph stats\n", + "print('\\n--- Graph Stats ---')\n", + "stats = graph_index.get_stats()\n", + "for key in ['source', 'chunk', 'topic', 'statement', 'fact', 'entity']:\n", + " print(f' {key:12s}: {stats.get(key, 0)}')\n", + "\n", + "print('\\n\\u2705 Pipeline completed!')\n", + "print(' NER and schema aligned on: Software, Format, Concept')\n", + "print(' Schema lightly filters out any entity types the LLM invented outside those three')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Available Stages Reference\n", + "\n", + "| Stage | Purpose | Requires |\n", + "|-------|---------|----------|\n", + "| `LLMPropositionStage` | Break text into atomic propositions via LLM | \u2014 |\n", + "| `LocalPropositionStage` | Break text into propositions locally (no LLM cost) | \u2014 |\n", + "| `LLMTopicExtractionStage` | Extract topics, entities, relationships via LLM | propositions (configurable) |\n", + "| `NERExtractionStage` | Find entities using local GLiNER model | `pip install gliner` |\n", + "| `EntityMergeStage` | Merge NER entities into LLM-extracted topics (optional fuzzy dedup) | NER output + topics |\n", + "| `SchemaFilterStage` | Remove entities/relationships not in schema | topics |\n", + "| `BatchLLMPropositionStage` | Batch proposition extraction via Bedrock | `BatchConfig` |\n", + "| `BatchTopicExtractionStage` | Batch topic extraction via Bedrock | `BatchConfig` |" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.6" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/lexical-graph/README.md b/lexical-graph/README.md index 5ae1e90c1..3bdee4a99 100644 --- a/lexical-graph/README.md +++ b/lexical-graph/README.md @@ -160,6 +160,7 @@ if __name__ == '__main__': - [Indexing](https://awslabs.github.io/graphrag-toolkit/lexical-graph/indexing/) - [Batch Extraction](https://awslabs.github.io/graphrag-toolkit/lexical-graph/batch-extraction/) - [Configuring Batch Extraction](https://awslabs.github.io/graphrag-toolkit/lexical-graph/configuring-batch-extraction/) + - [Composable Extraction Pipeline](https://awslabs.github.io/graphrag-toolkit/lexical-graph/composable-extraction-pipeline-usage/) - [Versioned Updates](https://awslabs.github.io/graphrag-toolkit/lexical-graph/versioned-updates/) - [Querying](https://awslabs.github.io/graphrag-toolkit/lexical-graph/querying/) - [Traversal-Based Search](https://awslabs.github.io/graphrag-toolkit/lexical-graph/traversal-based-search/) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/__init__.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/__init__.py index b984d341c..12b58e784 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/__init__.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/__init__.py @@ -12,3 +12,7 @@ from .infer_classifications import InferClassifications from .infer_config import InferClassificationsConfig from .preferred_values import PREFERRED_VALUES_PROVIDER_TYPE, PreferredValuesProvider, default_preferred_values +from .extraction_stage import ExtractionStage +from .extraction_schema import ExtractionSchema, EntityTypeConfig +from .pipeline_builder import PipelineBuilder +from .stages import LLMPropositionStage, LocalPropositionStage, LLMTopicExtractionStage, SchemaFilterStage, NERExtractionStage, EntityMergeStage, BatchLLMPropositionStage, BatchTopicExtractionStage diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/extraction_schema.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/extraction_schema.py new file mode 100644 index 000000000..49fe2269a --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/extraction_schema.py @@ -0,0 +1,60 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field +from typing import Dict, List, Optional + + +@dataclass +class EntityTypeConfig: + """Configuration for an entity type in the extraction schema. + + Attributes: + description: Human-readable description of this entity type. + attributes: Expected attribute names for this entity type. + aliases: Alternative names that should map to this entity type. + """ + description: Optional[str] = None + attributes: List[str] = field(default_factory=list) + aliases: List[str] = field(default_factory=list) + + +@dataclass +class ExtractionSchema: + """Schema defining allowed entity types, relationship types, and constraints. + + When provided to the extraction pipeline, the schema: + - Feeds entity type names into the extraction prompt as preferred classifications + - Configures NER stage labels (when using NERExtractionStage) + - Filters non-matching entities/relationships when strict=True + + Attributes: + entity_types: Mapping of entity type name to its configuration. + relationship_types: Allowed relationship type names. + strict: When True, filter out entities and relationships not in the schema. + """ + entity_types: Dict[str, EntityTypeConfig] = field(default_factory=dict) + relationship_types: List[str] = field(default_factory=list) + strict: bool = False + + def entity_type_names(self) -> List[str]: + """Return sorted list of entity type names.""" + return sorted(self.entity_types.keys()) + + def format_as_prompt_constraint(self) -> str: + """Format schema as a text block for injection into extraction prompts.""" + lines = [] + if self.entity_types: + lines.append('Entity types:') + for name, config in sorted(self.entity_types.items()): + desc = f' - {config.description}' if config.description else '' + lines.append(f' {name}{desc}') + if config.attributes: + lines.append(f' Attributes: {", ".join(config.attributes)}') + if self.relationship_types: + lines.append('Relationship types:') + for rt in sorted(self.relationship_types): + lines.append(f' {rt}') + if self.strict: + lines.append('STRICT MODE: Only extract entities and relationships matching the types listed above.') + return '\n'.join(lines) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/extraction_stage.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/extraction_stage.py new file mode 100644 index 000000000..e940f3bd9 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/extraction_stage.py @@ -0,0 +1,35 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from abc import ABC, abstractmethod +from typing import List, Optional + +from llama_index.core.schema import TransformComponent + + +class ExtractionStage(ABC): + """Abstract base class for composable extraction pipeline stages. + + Each stage declares its input and output metadata keys, enabling + the PipelineBuilder to validate stage compatibility at build time. + """ + + @abstractmethod + def input_keys(self) -> List[str]: + """Metadata keys this stage reads from.""" + ... + + @abstractmethod + def output_keys(self) -> List[str]: + """Metadata keys this stage writes to.""" + ... + + @abstractmethod + def as_transform(self) -> TransformComponent: + """Return the LlamaIndex TransformComponent for this stage.""" + ... + + @property + def stage_type(self) -> str: + """Stage type identifier: 'local', 'llm', 'filter', or 'transform'.""" + return 'transform' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/pipeline_builder.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/pipeline_builder.py new file mode 100644 index 000000000..89864e823 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/pipeline_builder.py @@ -0,0 +1,65 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +from typing import List, Optional + +from llama_index.core.schema import TransformComponent + +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage + +logger = logging.getLogger(__name__) + + +class PipelineBuilder: + """Builds a validated extraction pipeline from ExtractionStage instances. + + Validates that each stage's input_keys are satisfied by prior stages' + output_keys (or are available as initial keys). + """ + + def __init__(self, initial_keys: Optional[List[str]] = None): + self._stages: List[ExtractionStage] = [] + self._available_keys: set = set(initial_keys or []) + + def add(self, stage: ExtractionStage) -> 'PipelineBuilder': + """Add a stage to the pipeline with input key validation. + + Args: + stage: The extraction stage to add. + + Returns: + self for method chaining. + + Raises: + ValueError: If the stage requires input keys not yet available. + """ + missing = set(stage.input_keys()) - self._available_keys + if missing: + raise ValueError( + f"Stage {stage.__class__.__name__} requires keys {sorted(missing)} " + f"not provided by prior stages. Available: {sorted(self._available_keys)}" + ) + self._stages.append(stage) + self._available_keys.update(stage.output_keys()) + return self + + def build(self) -> List[TransformComponent]: + """Build the pipeline, returning a list of TransformComponents. + + Raises: + ValueError: If no stages have been added. + """ + if not self._stages: + raise ValueError('Pipeline has no stages') + return [stage.as_transform() for stage in self._stages] + + @property + def stages(self) -> List[ExtractionStage]: + """Return the current list of stages.""" + return list(self._stages) + + @property + def available_keys(self) -> List[str]: + """Return the currently available metadata keys.""" + return sorted(self._available_keys) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/__init__.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/__init__.py new file mode 100644 index 000000000..1063c35f6 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/__init__.py @@ -0,0 +1,10 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from .llm_proposition_stage import LLMPropositionStage +from .local_proposition_stage import LocalPropositionStage +from .llm_topic_extraction_stage import LLMTopicExtractionStage +from .schema_filter_stage import SchemaFilterStage +from .ner_extraction_stage import NERExtractionStage +from .entity_merge_stage import EntityMergeStage +from .batch_stages import BatchLLMPropositionStage, BatchTopicExtractionStage diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/batch_stages.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/batch_stages.py new file mode 100644 index 000000000..55c563feb --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/batch_stages.py @@ -0,0 +1,82 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import List, Optional + +from llama_index.core.schema import TransformComponent + +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.batch_llm_proposition_extractor_sync import BatchLLMPropositionExtractorSync +from graphrag_toolkit.lexical_graph.indexing.extract.batch_topic_extractor_sync import BatchTopicExtractorSync +from graphrag_toolkit.lexical_graph.indexing.extract.batch_config import BatchConfig +from graphrag_toolkit.lexical_graph.indexing.extract.preferred_values import PreferredValuesProvider +from graphrag_toolkit.lexical_graph.indexing.constants import PROPOSITIONS_KEY, TOPICS_KEY +from graphrag_toolkit.lexical_graph.utils.llm_cache import LLMCacheType + + +class BatchLLMPropositionStage(ExtractionStage): + """Stage wrapper for batch LLM proposition extraction via Bedrock.""" + + def __init__(self, batch_config: BatchConfig, prompt_template: Optional[str] = None, llm: Optional[LLMCacheType] = None): + self._batch_config = batch_config + self._prompt_template = prompt_template + self._llm = llm + + def input_keys(self) -> List[str]: + return [] + + def output_keys(self) -> List[str]: + return [PROPOSITIONS_KEY] + + def as_transform(self) -> TransformComponent: + return BatchLLMPropositionExtractorSync( + batch_config=self._batch_config, + prompt_template=self._prompt_template, + llm=self._llm, + ) + + @property + def stage_type(self) -> str: + return 'llm' + + +class BatchTopicExtractionStage(ExtractionStage): + """Stage wrapper for batch topic extraction via Bedrock.""" + + def __init__( + self, + batch_config: BatchConfig, + use_propositions: bool = True, + prompt_template: Optional[str] = None, + llm: Optional[LLMCacheType] = None, + entity_classification_provider: Optional[PreferredValuesProvider] = None, + topic_provider: Optional[PreferredValuesProvider] = None, + ): + self._batch_config = batch_config + self._use_propositions = use_propositions + self._prompt_template = prompt_template + self._llm = llm + self._entity_classification_provider = entity_classification_provider + self._topic_provider = topic_provider + + def input_keys(self) -> List[str]: + if self._use_propositions: + return [PROPOSITIONS_KEY] + return [] + + def output_keys(self) -> List[str]: + return [TOPICS_KEY] + + def as_transform(self) -> TransformComponent: + return BatchTopicExtractorSync( + batch_config=self._batch_config, + source_metadata_field=PROPOSITIONS_KEY if self._use_propositions else None, + prompt_template=self._prompt_template, + llm=self._llm, + entity_classification_provider=self._entity_classification_provider, + topic_provider=self._topic_provider, + ) + + @property + def stage_type(self) -> str: + return 'llm' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/entity_merge_stage.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/entity_merge_stage.py new file mode 100644 index 000000000..e29b26642 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/entity_merge_stage.py @@ -0,0 +1,92 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +from difflib import SequenceMatcher +from typing import List, Optional, Sequence + +from llama_index.core.schema import BaseNode, TransformComponent +from llama_index.core.bridge.pydantic import Field + +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.stages.ner_extraction_stage import PRE_EXTRACTED_ENTITIES_KEY +from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY +from graphrag_toolkit.lexical_graph.indexing.model import TopicCollection, Entity + +logger = logging.getLogger(__name__) + + +def _is_fuzzy_match(name: str, existing_names: set, threshold: float) -> bool: + """Check if name fuzzy-matches any existing name.""" + name_lower = name.lower() + for existing in existing_names: + if SequenceMatcher(None, name_lower, existing).ratio() >= threshold: + return True + return False + + +class EntityMergeTransform(TransformComponent): + """TransformComponent that merges NER-detected entities into TopicCollection.""" + + fuzzy_threshold: Optional[float] = Field( + default=None, + description='Fuzzy matching threshold (0-1). None means exact match only.', + ) + + def __call__(self, nodes: Sequence[BaseNode], **kwargs) -> Sequence[BaseNode]: + for node in nodes: + pre_extracted = node.metadata.get(PRE_EXTRACTED_ENTITIES_KEY, []) + topics_data = node.metadata.get(TOPICS_KEY) + + if not pre_extracted or not topics_data: + continue + + tc = TopicCollection(**topics_data) + + for topic in tc.topics: + existing_names = {e.value.lower() for e in topic.entities} + for ner_entity in pre_extracted: + name = ner_entity.get('value', '') + name_lower = name.lower() + + if self.fuzzy_threshold is not None: + is_dup = _is_fuzzy_match(name, existing_names, self.fuzzy_threshold) + else: + is_dup = name_lower in existing_names + + if not is_dup: + topic.entities.append(Entity( + value=name, + classification=ner_entity.get('classification', 'unknown'), + )) + existing_names.add(name_lower) + + node.metadata[TOPICS_KEY] = tc.model_dump() + + return nodes + + +class EntityMergeStage(ExtractionStage): + """Stage that merges NER-detected entities into LLM-extracted TopicCollection. + + Args: + fuzzy_threshold: Similarity threshold for fuzzy deduplication (0-1). + None (default) uses exact case-insensitive matching. + 0.85 is a good starting point for fuzzy matching. + """ + + def __init__(self, fuzzy_threshold: Optional[float] = None): + self._fuzzy_threshold = fuzzy_threshold + + def input_keys(self) -> List[str]: + return [PRE_EXTRACTED_ENTITIES_KEY, TOPICS_KEY] + + def output_keys(self) -> List[str]: + return [TOPICS_KEY] + + def as_transform(self) -> TransformComponent: + return EntityMergeTransform(fuzzy_threshold=self._fuzzy_threshold) + + @property + def stage_type(self) -> str: + return 'transform' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/llm_proposition_stage.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/llm_proposition_stage.py new file mode 100644 index 000000000..35691f110 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/llm_proposition_stage.py @@ -0,0 +1,35 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import List, Optional + +from llama_index.core.schema import TransformComponent + +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.llm_proposition_extractor import LLMPropositionExtractor +from graphrag_toolkit.lexical_graph.indexing.constants import PROPOSITIONS_KEY +from graphrag_toolkit.lexical_graph.utils.llm_cache import LLMCacheType + + +class LLMPropositionStage(ExtractionStage): + """Stage wrapper for LLM-based proposition extraction.""" + + def __init__(self, prompt_template: Optional[str] = None, llm: Optional[LLMCacheType] = None): + self._prompt_template = prompt_template + self._llm = llm + + def input_keys(self) -> List[str]: + return [] + + def output_keys(self) -> List[str]: + return [PROPOSITIONS_KEY] + + def as_transform(self) -> TransformComponent: + return LLMPropositionExtractor( + prompt_template=self._prompt_template, + llm=self._llm + ) + + @property + def stage_type(self) -> str: + return 'llm' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/llm_topic_extraction_stage.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/llm_topic_extraction_stage.py new file mode 100644 index 000000000..ef1940191 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/llm_topic_extraction_stage.py @@ -0,0 +1,59 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import List, Optional + +from llama_index.core.schema import TransformComponent + +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_schema import ExtractionSchema +from graphrag_toolkit.lexical_graph.indexing.extract.topic_extractor import TopicExtractor +from graphrag_toolkit.lexical_graph.indexing.extract.preferred_values import PreferredValuesProvider, default_preferred_values +from graphrag_toolkit.lexical_graph.indexing.constants import PROPOSITIONS_KEY, TOPICS_KEY +from graphrag_toolkit.lexical_graph.utils.llm_cache import LLMCacheType + + +class LLMTopicExtractionStage(ExtractionStage): + """Stage wrapper for LLM-based topic/entity/relationship extraction.""" + + def __init__( + self, + use_propositions: bool = True, + prompt_template: Optional[str] = None, + llm: Optional[LLMCacheType] = None, + entity_classification_provider: Optional[PreferredValuesProvider] = None, + topic_provider: Optional[PreferredValuesProvider] = None, + schema: Optional[ExtractionSchema] = None, + ): + self._use_propositions = use_propositions + self._prompt_template = prompt_template + self._llm = llm + self._entity_classification_provider = entity_classification_provider + self._topic_provider = topic_provider + self._schema = schema + + def input_keys(self) -> List[str]: + if self._use_propositions: + return [PROPOSITIONS_KEY] + return [] + + def output_keys(self) -> List[str]: + return [TOPICS_KEY] + + def as_transform(self) -> TransformComponent: + schema_constraints = '' + if self._schema: + schema_constraints = self._schema.format_as_prompt_constraint() + + return TopicExtractor( + source_metadata_field=PROPOSITIONS_KEY if self._use_propositions else None, + prompt_template=self._prompt_template, + llm=self._llm, + entity_classification_provider=self._entity_classification_provider, + topic_provider=self._topic_provider, + schema_constraints=schema_constraints, + ) + + @property + def stage_type(self) -> str: + return 'llm' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/local_proposition_stage.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/local_proposition_stage.py new file mode 100644 index 000000000..462680783 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/local_proposition_stage.py @@ -0,0 +1,36 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import List, Optional + +from llama_index.core.schema import TransformComponent + +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.proposition_extractor import PropositionExtractor +from graphrag_toolkit.lexical_graph.indexing.constants import PROPOSITIONS_KEY + + +class LocalPropositionStage(ExtractionStage): + """Stage wrapper for local transformer-based proposition extraction.""" + + def __init__(self, model_name: Optional[str] = None, device: Optional[str] = None): + self._model_name = model_name + self._device = device + + def input_keys(self) -> List[str]: + return [] + + def output_keys(self) -> List[str]: + return [PROPOSITIONS_KEY] + + def as_transform(self) -> TransformComponent: + kwargs = {} + if self._model_name: + kwargs['proposition_model_name'] = self._model_name + if self._device: + kwargs['device'] = self._device + return PropositionExtractor(**kwargs) + + @property + def stage_type(self) -> str: + return 'local' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/ner_extraction_stage.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/ner_extraction_stage.py new file mode 100644 index 000000000..67fa61978 --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/ner_extraction_stage.py @@ -0,0 +1,90 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +from typing import Any, List, Optional, Sequence + +from llama_index.core.schema import BaseNode, TransformComponent +from llama_index.core.bridge.pydantic import Field +from pydantic import PrivateAttr + +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage + +logger = logging.getLogger(__name__) + +PRE_EXTRACTED_ENTITIES_KEY = 'aws::graph::pre_extracted_entities' + + +class NERTransform(TransformComponent): + """TransformComponent that runs NER and writes entities to metadata.""" + + model_name: str = Field(default='urchade/gliner_base') + entity_labels: List[str] = Field(default_factory=list) + threshold: float = Field(default=0.5) + _model: Any = PrivateAttr(default=None) + + def __call__(self, nodes: Sequence[BaseNode], **kwargs) -> Sequence[BaseNode]: + try: + from gliner import GLiNER + except ImportError: + raise ImportError( + 'GLiNER is required for NERExtractionStage. ' + 'Install it with: pip install gliner' + ) + + if self._model is None: + logger.info('Loading NER model: %s', self.model_name) + try: + self._model = GLiNER.from_pretrained(self.model_name) + except Exception as e: + raise RuntimeError( + f'Failed to load NER model "{self.model_name}". ' + f'Ensure the model name is correct and you have internet access. ' + f'Error: {e}' + ) from e + + for node in nodes: + text = node.get_content() + if not text or not self.entity_labels: + node.metadata[PRE_EXTRACTED_ENTITIES_KEY] = [] + continue + + predictions = self._model.predict_entities(text, self.entity_labels, threshold=self.threshold) + entities = [ + {'value': p['text'], 'classification': p['label']} + for p in predictions + ] + node.metadata[PRE_EXTRACTED_ENTITIES_KEY] = entities + + return nodes + + +class NERExtractionStage(ExtractionStage): + """Stage wrapper for CPU-based NER extraction using GLiNER.""" + + def __init__( + self, + model_name: str = 'urchade/gliner_base', + entity_labels: Optional[List[str]] = None, + threshold: float = 0.5, + ): + self._model_name = model_name + self._entity_labels = entity_labels or [] + self._threshold = threshold + + def input_keys(self) -> List[str]: + return [] + + def output_keys(self) -> List[str]: + return [PRE_EXTRACTED_ENTITIES_KEY] + + def as_transform(self) -> TransformComponent: + return NERTransform( + model_name=self._model_name, + entity_labels=self._entity_labels, + threshold=self._threshold, + ) + + @property + def stage_type(self) -> str: + return 'local' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/schema_filter_stage.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/schema_filter_stage.py new file mode 100644 index 000000000..c88d7eb0b --- /dev/null +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/stages/schema_filter_stage.py @@ -0,0 +1,78 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import logging +from typing import List, Sequence + +from llama_index.core.schema import BaseNode, TransformComponent +from llama_index.core.bridge.pydantic import Field +from pydantic import ConfigDict + +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_schema import ExtractionSchema +from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY +from graphrag_toolkit.lexical_graph.indexing.model import TopicCollection + +logger = logging.getLogger(__name__) + + +class SchemaFilter(TransformComponent): + """TransformComponent that filters extracted topics against a schema.""" + + extraction_schema: ExtractionSchema = Field(description='Extraction schema for filtering') + + model_config = ConfigDict(arbitrary_types_allowed=True) + + def __call__(self, nodes: Sequence[BaseNode], **kwargs) -> Sequence[BaseNode]: + if not self.extraction_schema.strict: + return nodes + + allowed_entity_types = {n.lower() for n in self.extraction_schema.entity_type_names()} + for config in self.extraction_schema.entity_types.values(): + for alias in config.aliases: + allowed_entity_types.add(alias.lower()) + + allowed_relationships = {r.upper() for r in self.extraction_schema.relationship_types} if self.extraction_schema.relationship_types else None + + for node in nodes: + topics_data = node.metadata.get(TOPICS_KEY) + if not topics_data: + continue + + tc = TopicCollection(**topics_data) + for topic in tc.topics: + if allowed_entity_types: + topic.entities = [ + e for e in topic.entities + if e.classification and e.classification.lower() in allowed_entity_types + ] + if allowed_relationships: + for stmt in topic.statements: + stmt.facts = [ + f for f in stmt.facts + if f.predicate.value.upper() in allowed_relationships + ] + + node.metadata[TOPICS_KEY] = tc.model_dump() + + return nodes + + +class SchemaFilterStage(ExtractionStage): + """Stage that filters extracted topics against an ExtractionSchema.""" + + def __init__(self, schema: ExtractionSchema): + self._schema = schema + + def input_keys(self) -> List[str]: + return [TOPICS_KEY] + + def output_keys(self) -> List[str]: + return [TOPICS_KEY] + + def as_transform(self) -> TransformComponent: + return SchemaFilter(extraction_schema=self._schema) + + @property + def stage_type(self) -> str: + return 'filter' diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/topic_extractor.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/topic_extractor.py index efe3ceac3..033a61b8b 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/topic_extractor.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/topic_extractor.py @@ -7,7 +7,7 @@ from graphrag_toolkit.lexical_graph.config import GraphRAGConfig from graphrag_toolkit.lexical_graph.utils import LLMCache, LLMCacheType -from graphrag_toolkit.lexical_graph.indexing.utils.topic_utils import parse_extracted_topics, format_list, format_text +from graphrag_toolkit.lexical_graph.indexing.utils.topic_utils import parse_extracted_topics, parse_extracted_topics_json, format_list, format_text from graphrag_toolkit.lexical_graph.indexing.extract.preferred_values import PreferredValuesProvider, default_preferred_values from graphrag_toolkit.lexical_graph.indexing.model import TopicCollection from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY @@ -40,10 +40,20 @@ class TopicExtractor(BaseExtractor): description='Entity classification provider' ) + schema_constraints: str = Field( + default='', + description='Schema constraints to inject into the prompt' + ) + topic_provider:PreferredValuesProvider = Field( description='Topic provider' ) + output_format: str = Field( + default='text', + description="Output format: 'text' for legacy parsing, 'json' for structured JSON parsing with text fallback" + ) + @classmethod def class_name(cls) -> str: """ @@ -65,7 +75,9 @@ def __init__(self, source_metadata_field=None, num_workers:Optional[int]=None, entity_classification_provider=None, - topic_provider=None + topic_provider=None, + output_format:str='text', + schema_constraints:str='' ): """ Initializes the instance with the provided or default parameters to facilitate @@ -98,7 +110,9 @@ def __init__(self, source_metadata_field=source_metadata_field, num_workers=coalesce(num_workers, GraphRAGConfig.extraction_num_threads_per_worker), entity_classification_provider=entity_classification_provider or default_preferred_values([]), - topic_provider=topic_provider or default_preferred_values([]) + topic_provider=topic_provider or default_preferred_values([]), + output_format=output_format, + schema_constraints=schema_constraints ) logger.debug(f'Prompt template: {self.prompt_template}') @@ -205,6 +219,7 @@ def blocking_llm_call(): text=text, preferred_entity_classifications=format_list(preferred_entity_classifications), preferred_topics=format_list(preferred_topics), + schema_constraints=self.schema_constraints, #exclude_cache_keys=['preferred_entity_classifications', 'preferred_topics'] ) @@ -212,5 +227,11 @@ def blocking_llm_call(): raw_response = await coro - (topics, garbage) = parse_extracted_topics(raw_response) - return (topics, garbage) \ No newline at end of file + if self.output_format == 'json': + (topics, garbage) = parse_extracted_topics_json(raw_response) + if garbage: + logger.debug(f'JSON parsing failed, falling back to text parser: {garbage}') + (topics, garbage) = parse_extracted_topics(raw_response) + else: + (topics, garbage) = parse_extracted_topics(raw_response) + return (topics, garbage) diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py index 4534034b7..cc59a4ed9 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/prompts.py @@ -158,6 +158,8 @@ - Do not provide any other explanatory text - Extract only information explicitly stated in the propositions +{schema_constraints} + Adhere strictly to the provided instructions. Non-compliance will result in termination. @@ -230,3 +232,76 @@ Classification3 """ + +EXTRACT_TOPICS_JSON_PROMPT = """ +You are a top-tier algorithm designed for extracting information in structured formats to build a knowledge graph. +Your input consists of carefully crafted propositions - simple, atomic, and decontextualized statements. Your task is to: + 1. Organize these propositions into topics + 2. Extract entities and their attributes + 3. Identify relationships between entities + +Try to capture as much information from the text as possible without sacrificing accuracy. Do not add any information that is not explicitly mentioned in the input propositions. + +## Topic Extraction: + 1. Read the entire set of propositions and then extract a list of specific topics. Topic names should provide a clear, highly descriptive summary of the content. + - Compare each topic with the list of Preferred Topics. If a preferred topic matches the new topic in its meaning and specificity, use the preferred topic. + 2. Each proposition must be assigned to at least one topic. + +## Entity Extraction: + 1. Extract all named entities mentioned in the propositions within each topic. + 2. DO NOT treat numerical values, dates, times, measurements, or object attributes as entities. + 3. Classify each entity using the Preferred Entity Classifications when possible. + 4. Use the most complete identifier for each entity. Avoid articles at the beginning. + +## Relationship and Attribute Extraction: + 1. Extract relationships between entity pairs as: subject, predicate, object. + 2. Extract entity attributes as: subject, predicate (attribute name), complement (attribute value). + 3. Use general, timeless, uppercase relationship/attribute names with underscores. + +{schema_constraints} + +## Response Format: +Respond with a JSON object matching this exact schema: +```json +{{ + "topics": [ + {{ + "value": "topic name", + "entities": [ + {{"value": "entity name", "classification": "EntityType"}} + ], + "statements": [ + {{ + "value": "exact proposition text", + "facts": [ + {{ + "subject": {{"value": "entity name", "classification": "EntityType"}}, + "predicate": {{"value": "RELATIONSHIP_NAME"}}, + "object": {{"value": "entity name", "classification": "EntityType"}} + }} + ] + }} + ] + }} + ] +}} +``` + +## Strict Compliance: + - Use propositions exactly as provided + - Assign every proposition to at least one topic + - Respond ONLY with valid JSON, no other text + - Extract only information explicitly stated in the propositions + + +{text} + + + +{preferred_topics} + + + +{preferred_entity_classifications} + +""" diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/topic_utils.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/topic_utils.py index 6b7d6d512..96a4f1df4 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/topic_utils.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/utils/topic_utils.py @@ -44,6 +44,30 @@ def remove_articles(s:str): return s +def parse_extracted_topics_json(raw_text: str) -> Tuple[TopicCollection, List[str]]: + """Parse JSON-formatted topic extraction output into a TopicCollection. + + Returns: + Tuple of (TopicCollection, garbage list). Garbage is empty on success, + or contains the error message on failure. + """ + import json + + text = raw_text.strip() + # Strip markdown code fences if present + if text.startswith('```'): + lines = text.split('\n') + lines = [l for l in lines if not l.strip().startswith('```')] + text = '\n'.join(lines) + + try: + data = json.loads(text) + tc = TopicCollection(**data) + return (tc, []) + except Exception as e: + return (TopicCollection(topics=[]), [f'JSON_PARSE_ERROR: {e}']) + + def parse_extracted_topics(raw_text:str) -> Tuple[TopicCollection, List[str]]: garbage = [] diff --git a/lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py b/lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py index 6b60f4892..40a19170a 100644 --- a/lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py +++ b/lexical-graph/src/graphrag_toolkit/lexical_graph/lexical_graph_index.py @@ -25,6 +25,10 @@ from graphrag_toolkit.lexical_graph.indexing.extract import TopicExtractor, BatchTopicExtractorSync from graphrag_toolkit.lexical_graph.indexing.extract import ExtractionPipeline from graphrag_toolkit.lexical_graph.indexing.extract import InferClassifications, InferClassificationsConfig +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_schema import ExtractionSchema +from graphrag_toolkit.lexical_graph.indexing.extract.pipeline_builder import PipelineBuilder +from graphrag_toolkit.lexical_graph.indexing.extract.stages.llm_topic_extraction_stage import LLMTopicExtractionStage from graphrag_toolkit.lexical_graph.indexing.build import BuildPipeline from graphrag_toolkit.lexical_graph.indexing.build import VectorIndexing from graphrag_toolkit.lexical_graph.indexing.build import GraphConstruction @@ -53,6 +57,9 @@ class ExtractionConfig(): and topic or metadata filtering. It provides flexibility in customizing the extraction behavior and prompts to tailor it to specific use cases. + For custom composable pipelines, use the ExtractionConfig.from_stages() + class method instead of the constructor. + Attributes: enable_proposition_extraction (bool): Determines whether proposition extraction is enabled. Defaults to True. @@ -92,11 +99,32 @@ def __init__(self, self.extract_propositions_prompt_template = extract_propositions_prompt_template self.extract_topics_prompt_template = extract_topics_prompt_template self.extraction_filters = FilterConfig(extraction_filters) + self.stages = None + self.schema = None if extraction_llm is not None: self.extraction_llm = extraction_llm if isinstance(extraction_llm, LLMCache) else GraphRAGConfig.to_llm(extraction_llm) else: self.extraction_llm = None + @classmethod + def from_stages(cls, stages: List['ExtractionStage'], schema: Optional['ExtractionSchema'] = None): + """Create config with a custom composable extraction pipeline. + + When using custom stages, configure extraction behavior (LLM, prompts, + classifications) directly on the stage instances. + + Args: + stages: Ordered list of ExtractionStage instances defining the pipeline. + schema: Optional ExtractionSchema for type constraints and filtering. + + Returns: + ExtractionConfig configured for custom pipeline execution. + """ + config = cls() + config.stages = stages + config.schema = schema + return config + class BuildConfig(): """ @@ -342,6 +370,22 @@ def _configure_extraction_pipeline(self, config: IndexingConfig): for c in config.chunking: components.append(c) + # If custom stages are provided, use PipelineBuilder + if config.extraction.stages: + schema = config.extraction.schema + builder = PipelineBuilder() + for stage in config.extraction.stages: + # Auto-inject schema into LLMTopicExtractionStage if not already set + if schema and isinstance(stage, LLMTopicExtractionStage) and stage._schema is None: + stage._schema = schema + builder.add(stage) + components.extend(builder.build()) + logger.info( + 'Custom extraction pipeline: %s', + ' → '.join(type(s).__name__ for s in config.extraction.stages) + ) + return (pre_processors, components) + if config.extraction.enable_proposition_extraction: if config.batch_config: components.append(BatchLLMPropositionExtractorSync( diff --git a/lexical-graph/tests/unit/indexing/extract/test_batch_stages.py b/lexical-graph/tests/unit/indexing/extract/test_batch_stages.py new file mode 100644 index 000000000..0daaf7ea1 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_batch_stages.py @@ -0,0 +1,64 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from unittest.mock import MagicMock +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.stages.batch_stages import ( + BatchLLMPropositionStage, + BatchTopicExtractionStage, +) +from graphrag_toolkit.lexical_graph.indexing.constants import PROPOSITIONS_KEY, TOPICS_KEY + + +class TestBatchLLMPropositionStage: + """Tests for BatchLLMPropositionStage.""" + + def test_is_extraction_stage(self): + """Verify BatchLLMPropositionStage is an ExtractionStage.""" + stage = BatchLLMPropositionStage(batch_config=MagicMock()) + assert isinstance(stage, ExtractionStage) + + def test_input_keys_empty(self): + """Verify input_keys is empty.""" + stage = BatchLLMPropositionStage(batch_config=MagicMock()) + assert stage.input_keys() == [] + + def test_output_keys(self): + """Verify output_keys returns PROPOSITIONS_KEY.""" + stage = BatchLLMPropositionStage(batch_config=MagicMock()) + assert stage.output_keys() == [PROPOSITIONS_KEY] + + def test_stage_type(self): + """Verify stage_type is 'llm'.""" + stage = BatchLLMPropositionStage(batch_config=MagicMock()) + assert stage.stage_type == 'llm' + + +class TestBatchTopicExtractionStage: + """Tests for BatchTopicExtractionStage.""" + + def test_is_extraction_stage(self): + """Verify BatchTopicExtractionStage is an ExtractionStage.""" + stage = BatchTopicExtractionStage(batch_config=MagicMock()) + assert isinstance(stage, ExtractionStage) + + def test_input_keys_with_propositions(self): + """Verify input_keys includes PROPOSITIONS_KEY by default.""" + stage = BatchTopicExtractionStage(batch_config=MagicMock()) + assert stage.input_keys() == [PROPOSITIONS_KEY] + + def test_input_keys_without_propositions(self): + """Verify input_keys is empty when use_propositions=False.""" + stage = BatchTopicExtractionStage(batch_config=MagicMock(), use_propositions=False) + assert stage.input_keys() == [] + + def test_output_keys(self): + """Verify output_keys returns TOPICS_KEY.""" + stage = BatchTopicExtractionStage(batch_config=MagicMock()) + assert stage.output_keys() == [TOPICS_KEY] + + def test_stage_type(self): + """Verify stage_type is 'llm'.""" + stage = BatchTopicExtractionStage(batch_config=MagicMock()) + assert stage.stage_type == 'llm' diff --git a/lexical-graph/tests/unit/indexing/extract/test_extraction_schema.py b/lexical-graph/tests/unit/indexing/extract/test_extraction_schema.py new file mode 100644 index 000000000..0da022527 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_extraction_schema.py @@ -0,0 +1,105 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_schema import ( + EntityTypeConfig, + ExtractionSchema, +) + + +class TestEntityTypeConfig: + """Tests for EntityTypeConfig.""" + + def test_default_values(self): + """Verify defaults are None/empty.""" + config = EntityTypeConfig() + assert config.description is None + assert config.attributes == [] + assert config.aliases == [] + + def test_with_all_fields(self): + """Verify all fields can be set.""" + config = EntityTypeConfig( + description='A person entity', + attributes=['name', 'age'], + aliases=['Individual', 'Human'], + ) + assert config.description == 'A person entity' + assert config.attributes == ['name', 'age'] + assert config.aliases == ['Individual', 'Human'] + + +class TestExtractionSchema: + """Tests for ExtractionSchema.""" + + def test_default_values(self): + """Verify defaults are empty/False.""" + schema = ExtractionSchema() + assert schema.entity_types == {} + assert schema.relationship_types == [] + assert schema.strict is False + + def test_entity_type_names(self): + """Verify entity_type_names returns sorted names.""" + schema = ExtractionSchema(entity_types={ + 'Person': EntityTypeConfig(), + 'Company': EntityTypeConfig(), + 'Location': EntityTypeConfig(), + }) + assert schema.entity_type_names() == ['Company', 'Location', 'Person'] + + def test_entity_type_names_empty(self): + """Verify entity_type_names returns empty list when no types.""" + schema = ExtractionSchema() + assert schema.entity_type_names() == [] + + def test_format_as_prompt_constraint_empty(self): + """Verify empty schema produces empty string.""" + schema = ExtractionSchema() + assert schema.format_as_prompt_constraint() == '' + + def test_format_as_prompt_constraint_entity_types(self): + """Verify entity types are formatted correctly.""" + schema = ExtractionSchema(entity_types={ + 'Person': EntityTypeConfig(description='A human being', attributes=['name', 'age']), + }) + result = schema.format_as_prompt_constraint() + assert 'Entity types:' in result + assert 'Person - A human being' in result + assert 'Attributes: name, age' in result + + def test_format_as_prompt_constraint_relationship_types(self): + """Verify relationship types are formatted correctly.""" + schema = ExtractionSchema(relationship_types=['WORKS_FOR', 'LOCATED_IN']) + result = schema.format_as_prompt_constraint() + assert 'Relationship types:' in result + assert 'LOCATED_IN' in result + assert 'WORKS_FOR' in result + + def test_format_as_prompt_constraint_strict(self): + """Verify strict mode adds constraint text.""" + schema = ExtractionSchema( + entity_types={'Person': EntityTypeConfig()}, + strict=True, + ) + result = schema.format_as_prompt_constraint() + assert 'STRICT MODE' in result + + def test_format_as_prompt_constraint_not_strict(self): + """Verify non-strict mode omits constraint text.""" + schema = ExtractionSchema( + entity_types={'Person': EntityTypeConfig()}, + strict=False, + ) + result = schema.format_as_prompt_constraint() + assert 'STRICT MODE' not in result + + def test_format_as_prompt_constraint_no_description(self): + """Verify entity without description omits dash.""" + schema = ExtractionSchema(entity_types={ + 'Person': EntityTypeConfig(), + }) + result = schema.format_as_prompt_constraint() + assert 'Person' in result + assert ' - ' not in result diff --git a/lexical-graph/tests/unit/indexing/extract/test_extraction_stage.py b/lexical-graph/tests/unit/indexing/extract/test_extraction_stage.py new file mode 100644 index 000000000..f10d6fbfd --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_extraction_stage.py @@ -0,0 +1,64 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from unittest.mock import MagicMock +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage + + +class ConcreteStage(ExtractionStage): + """Concrete implementation for testing.""" + + def input_keys(self): + return ['input_key'] + + def output_keys(self): + return ['output_key'] + + def as_transform(self): + return MagicMock() + + +class TestExtractionStage: + """Tests for ExtractionStage ABC.""" + + def test_cannot_instantiate_abstract_class(self): + """Verify ExtractionStage cannot be instantiated directly.""" + with pytest.raises(TypeError): + ExtractionStage() + + def test_concrete_stage_instantiation(self): + """Verify concrete implementation can be instantiated.""" + stage = ConcreteStage() + assert stage is not None + + def test_input_keys(self): + """Verify input_keys returns expected keys.""" + stage = ConcreteStage() + assert stage.input_keys() == ['input_key'] + + def test_output_keys(self): + """Verify output_keys returns expected keys.""" + stage = ConcreteStage() + assert stage.output_keys() == ['output_key'] + + def test_as_transform_returns_component(self): + """Verify as_transform returns a TransformComponent.""" + stage = ConcreteStage() + result = stage.as_transform() + assert result is not None + + def test_default_stage_type(self): + """Verify default stage_type is 'transform'.""" + stage = ConcreteStage() + assert stage.stage_type == 'transform' + + def test_custom_stage_type(self): + """Verify stage_type can be overridden.""" + class LLMStage(ConcreteStage): + @property + def stage_type(self): + return 'llm' + + stage = LLMStage() + assert stage.stage_type == 'llm' diff --git a/lexical-graph/tests/unit/indexing/extract/test_json_topic_extraction.py b/lexical-graph/tests/unit/indexing/extract/test_json_topic_extraction.py new file mode 100644 index 000000000..22613d127 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_json_topic_extraction.py @@ -0,0 +1,86 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import json +from graphrag_toolkit.lexical_graph.indexing.utils.topic_utils import parse_extracted_topics_json +from graphrag_toolkit.lexical_graph.indexing.model import TopicCollection +from graphrag_toolkit.lexical_graph.indexing.prompts import EXTRACT_TOPICS_JSON_PROMPT + + +class TestParseExtractedTopicsJson: + """Tests for parse_extracted_topics_json function.""" + + def test_valid_json(self): + """Verify valid JSON is parsed into TopicCollection.""" + data = { + "topics": [{ + "value": "Test Topic", + "entities": [{"value": "John", "classification": "Person"}], + "statements": [{ + "value": "John works at Acme", + "facts": [{ + "subject": {"value": "John", "classification": "Person"}, + "predicate": {"value": "WORKS_FOR"}, + "object": {"value": "Acme", "classification": "Company"}, + }] + }] + }] + } + tc, garbage = parse_extracted_topics_json(json.dumps(data)) + assert len(garbage) == 0 + assert len(tc.topics) == 1 + assert tc.topics[0].value == "Test Topic" + assert len(tc.topics[0].entities) == 1 + assert tc.topics[0].entities[0].value == "John" + assert len(tc.topics[0].statements[0].facts) == 1 + + def test_json_with_code_fences(self): + """Verify JSON wrapped in markdown code fences is parsed.""" + data = {"topics": [{"value": "Topic", "entities": [], "statements": []}]} + raw = f"```json\n{json.dumps(data)}\n```" + tc, garbage = parse_extracted_topics_json(raw) + assert len(garbage) == 0 + assert len(tc.topics) == 1 + + def test_empty_topics(self): + """Verify empty topics list is valid.""" + tc, garbage = parse_extracted_topics_json('{"topics": []}') + assert len(garbage) == 0 + assert len(tc.topics) == 0 + + def test_invalid_json_returns_error(self): + """Verify invalid JSON returns empty TopicCollection with error.""" + tc, garbage = parse_extracted_topics_json('not json at all') + assert len(tc.topics) == 0 + assert len(garbage) == 1 + assert 'JSON_PARSE_ERROR' in garbage[0] + + def test_malformed_structure_returns_error(self): + """Verify JSON with wrong structure returns error.""" + tc, garbage = parse_extracted_topics_json('{"wrong_key": []}') + assert len(tc.topics) == 0 + + def test_whitespace_handling(self): + """Verify leading/trailing whitespace is handled.""" + data = {"topics": [{"value": "Topic", "entities": [], "statements": []}]} + raw = f" \n{json.dumps(data)}\n " + tc, garbage = parse_extracted_topics_json(raw) + assert len(garbage) == 0 + assert len(tc.topics) == 1 + + +class TestExtractTopicsJsonPrompt: + """Tests for EXTRACT_TOPICS_JSON_PROMPT template.""" + + def test_prompt_has_required_placeholders(self): + """Verify prompt contains all required template variables.""" + assert '{text}' in EXTRACT_TOPICS_JSON_PROMPT + assert '{preferred_topics}' in EXTRACT_TOPICS_JSON_PROMPT + assert '{preferred_entity_classifications}' in EXTRACT_TOPICS_JSON_PROMPT + assert '{schema_constraints}' in EXTRACT_TOPICS_JSON_PROMPT + + def test_prompt_requests_json_output(self): + """Verify prompt instructs JSON output.""" + assert 'JSON' in EXTRACT_TOPICS_JSON_PROMPT + assert '"topics"' in EXTRACT_TOPICS_JSON_PROMPT diff --git a/lexical-graph/tests/unit/indexing/extract/test_ner_and_merge_stages.py b/lexical-graph/tests/unit/indexing/extract/test_ner_and_merge_stages.py new file mode 100644 index 000000000..9ac4a4689 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_ner_and_merge_stages.py @@ -0,0 +1,175 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from unittest.mock import patch, MagicMock +from llama_index.core.schema import TextNode +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.stages.ner_extraction_stage import ( + NERExtractionStage, NERTransform, PRE_EXTRACTED_ENTITIES_KEY, +) +from graphrag_toolkit.lexical_graph.indexing.extract.stages.entity_merge_stage import ( + EntityMergeStage, EntityMergeTransform, +) +from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY +from graphrag_toolkit.lexical_graph.indexing.model import ( + TopicCollection, Topic, Entity, +) + + +class TestNERExtractionStage: + """Tests for NERExtractionStage.""" + + def test_is_extraction_stage(self): + stage = NERExtractionStage() + assert isinstance(stage, ExtractionStage) + + def test_input_keys_empty(self): + stage = NERExtractionStage() + assert stage.input_keys() == [] + + def test_output_keys(self): + stage = NERExtractionStage() + assert stage.output_keys() == [PRE_EXTRACTED_ENTITIES_KEY] + + def test_stage_type(self): + stage = NERExtractionStage() + assert stage.stage_type == 'local' + + def test_as_transform_returns_ner_transform(self): + stage = NERExtractionStage(model_name='test-model', entity_labels=['Person'], threshold=0.7) + transform = stage.as_transform() + assert isinstance(transform, NERTransform) + assert transform.model_name == 'test-model' + assert transform.entity_labels == ['Person'] + assert transform.threshold == 0.7 + + +class TestEntityMergeStage: + """Tests for EntityMergeStage.""" + + def test_is_extraction_stage(self): + stage = EntityMergeStage() + assert isinstance(stage, ExtractionStage) + + def test_input_keys(self): + stage = EntityMergeStage() + assert PRE_EXTRACTED_ENTITIES_KEY in stage.input_keys() + assert TOPICS_KEY in stage.input_keys() + + def test_output_keys(self): + stage = EntityMergeStage() + assert stage.output_keys() == [TOPICS_KEY] + + def test_stage_type(self): + stage = EntityMergeStage() + assert stage.stage_type == 'transform' + + +class TestEntityMergeTransform: + """Tests for EntityMergeTransform.""" + + def _make_node(self, ner_entities, topic_entities): + tc = TopicCollection(topics=[ + Topic(value='test', entities=topic_entities, statements=[]) + ]) + node = TextNode(text='test') + node.metadata[PRE_EXTRACTED_ENTITIES_KEY] = ner_entities + node.metadata[TOPICS_KEY] = tc.model_dump() + return node + + def test_merges_new_entities(self): + """Verify NER entities not in topics are added.""" + ner = [{'value': 'Acme', 'classification': 'Company'}] + existing = [Entity(value='John', classification='Person')] + node = self._make_node(ner, existing) + + transform = EntityMergeTransform() + result = transform([node]) + + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + names = {e.value for e in tc.topics[0].entities} + assert 'John' in names + assert 'Acme' in names + + def test_skips_duplicate_entities(self): + """Verify NER entities already in topics are not duplicated.""" + ner = [{'value': 'John', 'classification': 'Person'}] + existing = [Entity(value='John', classification='Person')] + node = self._make_node(ner, existing) + + transform = EntityMergeTransform() + result = transform([node]) + + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].entities) == 1 + + def test_case_insensitive_dedup(self): + """Verify deduplication is case-insensitive.""" + ner = [{'value': 'john', 'classification': 'Person'}] + existing = [Entity(value='John', classification='Person')] + node = self._make_node(ner, existing) + + transform = EntityMergeTransform() + result = transform([node]) + + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].entities) == 1 + + def test_no_ner_entities_unchanged(self): + """Verify node without NER entities is unchanged.""" + existing = [Entity(value='John', classification='Person')] + tc = TopicCollection(topics=[Topic(value='test', entities=existing, statements=[])]) + node = TextNode(text='test') + node.metadata[TOPICS_KEY] = tc.model_dump() + + transform = EntityMergeTransform() + result = transform([node]) + + tc2 = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc2.topics[0].entities) == 1 + + def test_no_topics_unchanged(self): + """Verify node without topics is unchanged.""" + node = TextNode(text='test') + node.metadata[PRE_EXTRACTED_ENTITIES_KEY] = [{'value': 'X', 'classification': 'Y'}] + + transform = EntityMergeTransform() + result = transform([node]) + assert TOPICS_KEY not in result[0].metadata + + def test_fuzzy_dedup_blocks_similar(self): + """Verify fuzzy matching blocks similar entity names.""" + ner = [{'value': 'DataBridge', 'classification': 'Company'}] + existing = [Entity(value='DataBridge AI', classification='Company')] + node = self._make_node(ner, existing) + + transform = EntityMergeTransform(fuzzy_threshold=0.7) + result = transform([node]) + + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].entities) == 1 + + def test_fuzzy_dedup_allows_different(self): + """Verify fuzzy matching allows sufficiently different names.""" + ner = [{'value': 'Acme Corp', 'classification': 'Company'}] + existing = [Entity(value='DataBridge AI', classification='Company')] + node = self._make_node(ner, existing) + + transform = EntityMergeTransform(fuzzy_threshold=0.7) + result = transform([node]) + + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].entities) == 2 + + def test_no_fuzzy_threshold_exact_match(self): + """Verify None threshold uses exact matching (similar names not deduped).""" + ner = [{'value': 'DataBridge', 'classification': 'Company'}] + existing = [Entity(value='DataBridge AI', classification='Company')] + node = self._make_node(ner, existing) + + transform = EntityMergeTransform(fuzzy_threshold=None) + result = transform([node]) + + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].entities) == 2 diff --git a/lexical-graph/tests/unit/indexing/extract/test_pipeline_builder.py b/lexical-graph/tests/unit/indexing/extract/test_pipeline_builder.py new file mode 100644 index 000000000..380f04267 --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_pipeline_builder.py @@ -0,0 +1,99 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from unittest.mock import MagicMock +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.pipeline_builder import PipelineBuilder + + +def make_stage(input_keys, output_keys): + """Helper to create a concrete ExtractionStage.""" + class TestStage(ExtractionStage): + def __init__(self, ik, ok): + self._ik = ik + self._ok = ok + def input_keys(self): + return self._ik + def output_keys(self): + return self._ok + def as_transform(self): + return MagicMock() + return TestStage(input_keys, output_keys) + + +class TestPipelineBuilder: + """Tests for PipelineBuilder.""" + + def test_build_empty_raises(self): + """Verify building an empty pipeline raises ValueError.""" + builder = PipelineBuilder() + with pytest.raises(ValueError, match='no stages'): + builder.build() + + def test_add_stage_no_inputs(self): + """Verify a stage with no input requirements can be added.""" + builder = PipelineBuilder() + stage = make_stage([], ['key_a']) + builder.add(stage) + assert len(builder.stages) == 1 + + def test_add_stage_with_satisfied_inputs(self): + """Verify a stage with satisfied inputs can be added.""" + builder = PipelineBuilder(initial_keys=['key_a']) + stage = make_stage(['key_a'], ['key_b']) + builder.add(stage) + assert len(builder.stages) == 1 + + def test_add_stage_with_missing_inputs_raises(self): + """Verify adding a stage with unsatisfied inputs raises ValueError.""" + builder = PipelineBuilder() + stage = make_stage(['missing_key'], ['key_b']) + with pytest.raises(ValueError, match='missing_key'): + builder.add(stage) + + def test_chained_stages(self): + """Verify stages can be chained where output feeds input.""" + builder = PipelineBuilder() + stage1 = make_stage([], ['key_a']) + stage2 = make_stage(['key_a'], ['key_b']) + stage3 = make_stage(['key_a', 'key_b'], ['key_c']) + builder.add(stage1).add(stage2).add(stage3) + assert len(builder.stages) == 3 + + def test_build_returns_transforms(self): + """Verify build returns list of TransformComponents.""" + builder = PipelineBuilder() + stage = make_stage([], ['key_a']) + builder.add(stage) + result = builder.build() + assert len(result) == 1 + + def test_available_keys_tracks_outputs(self): + """Verify available_keys accumulates output keys.""" + builder = PipelineBuilder(initial_keys=['init']) + stage = make_stage([], ['key_a', 'key_b']) + builder.add(stage) + assert 'init' in builder.available_keys + assert 'key_a' in builder.available_keys + assert 'key_b' in builder.available_keys + + def test_method_chaining(self): + """Verify add() returns self for chaining.""" + builder = PipelineBuilder() + result = builder.add(make_stage([], ['a'])) + assert result is builder + + def test_initial_keys_satisfy_first_stage(self): + """Verify initial_keys are available for the first stage.""" + builder = PipelineBuilder(initial_keys=['text', 'source']) + stage = make_stage(['text', 'source'], ['propositions']) + builder.add(stage) + assert len(builder.stages) == 1 + + def test_partial_missing_keys_raises(self): + """Verify error when only some required keys are available.""" + builder = PipelineBuilder(initial_keys=['key_a']) + stage = make_stage(['key_a', 'key_b'], ['key_c']) + with pytest.raises(ValueError, match='key_b'): + builder.add(stage) diff --git a/lexical-graph/tests/unit/indexing/extract/test_schema_filter_stage.py b/lexical-graph/tests/unit/indexing/extract/test_schema_filter_stage.py new file mode 100644 index 000000000..de25e74fb --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_schema_filter_stage.py @@ -0,0 +1,198 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from llama_index.core.schema import TextNode +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_schema import ExtractionSchema, EntityTypeConfig +from graphrag_toolkit.lexical_graph.indexing.extract.stages.schema_filter_stage import SchemaFilter, SchemaFilterStage +from graphrag_toolkit.lexical_graph.indexing.constants import TOPICS_KEY +from graphrag_toolkit.lexical_graph.indexing.model import ( + TopicCollection, Topic, Statement, Fact, Entity, Relation, +) + + +def make_topic_node(topics_data): + """Create a TextNode with topics metadata.""" + node = TextNode(text='test') + node.metadata[TOPICS_KEY] = topics_data + return node + + +def make_topics_with_entities(entities, facts=None): + """Create a TopicCollection dict with given entities and optional facts.""" + stmts = [] + if facts: + stmts = [Statement(value='test stmt', facts=facts)] + tc = TopicCollection(topics=[ + Topic(value='test topic', entities=entities, statements=stmts) + ]) + return tc.model_dump() + + +class TestSchemaFilterStage: + """Tests for SchemaFilterStage.""" + + def test_is_extraction_stage(self): + """Verify SchemaFilterStage is an ExtractionStage.""" + stage = SchemaFilterStage(schema=ExtractionSchema()) + assert isinstance(stage, ExtractionStage) + + def test_input_keys(self): + """Verify input_keys returns TOPICS_KEY.""" + stage = SchemaFilterStage(schema=ExtractionSchema()) + assert stage.input_keys() == [TOPICS_KEY] + + def test_output_keys(self): + """Verify output_keys returns TOPICS_KEY.""" + stage = SchemaFilterStage(schema=ExtractionSchema()) + assert stage.output_keys() == [TOPICS_KEY] + + def test_stage_type(self): + """Verify stage_type is 'filter'.""" + stage = SchemaFilterStage(schema=ExtractionSchema()) + assert stage.stage_type == 'filter' + + def test_as_transform_returns_schema_filter(self): + """Verify as_transform returns a SchemaFilter.""" + stage = SchemaFilterStage(schema=ExtractionSchema()) + result = stage.as_transform() + assert isinstance(result, SchemaFilter) + + +class TestSchemaFilter: + """Tests for SchemaFilter TransformComponent.""" + + def test_non_strict_passes_through(self): + """Verify non-strict mode passes all data through unchanged.""" + schema = ExtractionSchema( + entity_types={'Person': EntityTypeConfig()}, + strict=False, + ) + entities = [ + Entity(value='John', classification='Person'), + Entity(value='Acme', classification='Company'), + ] + node = make_topic_node(make_topics_with_entities(entities)) + sf = SchemaFilter(extraction_schema=schema) + result = sf([node]) + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].entities) == 2 + + def test_strict_filters_entities(self): + """Verify strict mode removes entities not in schema.""" + schema = ExtractionSchema( + entity_types={'Person': EntityTypeConfig()}, + strict=True, + ) + entities = [ + Entity(value='John', classification='Person'), + Entity(value='Acme', classification='Company'), + ] + node = make_topic_node(make_topics_with_entities(entities)) + sf = SchemaFilter(extraction_schema=schema) + result = sf([node]) + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].entities) == 1 + assert tc.topics[0].entities[0].value == 'John' + + def test_strict_filters_relationships(self): + """Verify strict mode removes facts with disallowed relationship types.""" + schema = ExtractionSchema( + entity_types={ + 'Person': EntityTypeConfig(), + 'Company': EntityTypeConfig(), + }, + relationship_types=['WORKS_FOR'], + strict=True, + ) + facts = [ + Fact( + subject=Entity(value='John', classification='Person'), + predicate=Relation(value='WORKS_FOR'), + object=Entity(value='Acme', classification='Company'), + ), + Fact( + subject=Entity(value='John', classification='Person'), + predicate=Relation(value='LIVES_IN'), + object=Entity(value='NYC', classification='Location'), + ), + ] + entities = [ + Entity(value='John', classification='Person'), + Entity(value='Acme', classification='Company'), + Entity(value='NYC', classification='Location'), + ] + node = make_topic_node(make_topics_with_entities(entities, facts)) + sf = SchemaFilter(extraction_schema=schema) + result = sf([node]) + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].statements[0].facts) == 1 + assert tc.topics[0].statements[0].facts[0].predicate.value == 'WORKS_FOR' + + def test_strict_entity_alias_matching(self): + """Verify strict mode matches entity aliases.""" + schema = ExtractionSchema( + entity_types={ + 'Person': EntityTypeConfig(aliases=['Individual']), + }, + strict=True, + ) + entities = [ + Entity(value='John', classification='Individual'), + Entity(value='Acme', classification='Company'), + ] + node = make_topic_node(make_topics_with_entities(entities)) + sf = SchemaFilter(extraction_schema=schema) + result = sf([node]) + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].entities) == 1 + assert tc.topics[0].entities[0].value == 'John' + + def test_strict_case_insensitive_entity_matching(self): + """Verify entity type matching is case-insensitive.""" + schema = ExtractionSchema( + entity_types={'Person': EntityTypeConfig()}, + strict=True, + ) + entities = [ + Entity(value='John', classification='person'), + Entity(value='Jane', classification='PERSON'), + ] + node = make_topic_node(make_topics_with_entities(entities)) + sf = SchemaFilter(extraction_schema=schema) + result = sf([node]) + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].entities) == 2 + + def test_node_without_topics_unchanged(self): + """Verify nodes without topics metadata are unchanged.""" + schema = ExtractionSchema(entity_types={'Person': EntityTypeConfig()}, strict=True) + node = TextNode(text='test') + sf = SchemaFilter(extraction_schema=schema) + result = sf([node]) + assert TOPICS_KEY not in result[0].metadata + + def test_no_relationship_types_skips_fact_filtering(self): + """Verify no relationship_types means facts are not filtered.""" + schema = ExtractionSchema( + entity_types={'Person': EntityTypeConfig(), 'Company': EntityTypeConfig()}, + relationship_types=[], + strict=True, + ) + facts = [ + Fact( + subject=Entity(value='John', classification='Person'), + predicate=Relation(value='WORKS_FOR'), + object=Entity(value='Acme', classification='Company'), + ), + ] + entities = [ + Entity(value='John', classification='Person'), + Entity(value='Acme', classification='Company'), + ] + node = make_topic_node(make_topics_with_entities(entities, facts)) + sf = SchemaFilter(extraction_schema=schema) + result = sf([node]) + tc = TopicCollection(**result[0].metadata[TOPICS_KEY]) + assert len(tc.topics[0].statements[0].facts) == 1 diff --git a/lexical-graph/tests/unit/indexing/extract/test_stages.py b/lexical-graph/tests/unit/indexing/extract/test_stages.py new file mode 100644 index 000000000..e433aae9b --- /dev/null +++ b/lexical-graph/tests/unit/indexing/extract/test_stages.py @@ -0,0 +1,138 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +from unittest.mock import patch, MagicMock +from graphrag_toolkit.lexical_graph.indexing.extract.extraction_stage import ExtractionStage +from graphrag_toolkit.lexical_graph.indexing.extract.stages import ( + LLMPropositionStage, + LocalPropositionStage, + LLMTopicExtractionStage, +) +from graphrag_toolkit.lexical_graph.indexing.constants import PROPOSITIONS_KEY, TOPICS_KEY + + +class TestLLMPropositionStage: + """Tests for LLMPropositionStage.""" + + def test_is_extraction_stage(self): + """Verify LLMPropositionStage is an ExtractionStage.""" + stage = LLMPropositionStage() + assert isinstance(stage, ExtractionStage) + + def test_input_keys_empty(self): + """Verify input_keys returns empty list (reads raw text).""" + stage = LLMPropositionStage() + assert stage.input_keys() == [] + + def test_output_keys(self): + """Verify output_keys returns PROPOSITIONS_KEY.""" + stage = LLMPropositionStage() + assert stage.output_keys() == [PROPOSITIONS_KEY] + + def test_stage_type(self): + """Verify stage_type is 'llm'.""" + stage = LLMPropositionStage() + assert stage.stage_type == 'llm' + + @patch('graphrag_toolkit.lexical_graph.indexing.extract.stages.llm_proposition_stage.LLMPropositionExtractor') + def test_as_transform_creates_extractor(self, mock_cls): + """Verify as_transform creates LLMPropositionExtractor.""" + stage = LLMPropositionStage(prompt_template='test prompt') + stage.as_transform() + mock_cls.assert_called_once_with(prompt_template='test prompt', llm=None) + + +class TestLocalPropositionStage: + """Tests for LocalPropositionStage.""" + + def test_is_extraction_stage(self): + """Verify LocalPropositionStage is an ExtractionStage.""" + stage = LocalPropositionStage() + assert isinstance(stage, ExtractionStage) + + def test_input_keys_empty(self): + """Verify input_keys returns empty list.""" + stage = LocalPropositionStage() + assert stage.input_keys() == [] + + def test_output_keys(self): + """Verify output_keys returns PROPOSITIONS_KEY.""" + stage = LocalPropositionStage() + assert stage.output_keys() == [PROPOSITIONS_KEY] + + def test_stage_type(self): + """Verify stage_type is 'local'.""" + stage = LocalPropositionStage() + assert stage.stage_type == 'local' + + @patch('graphrag_toolkit.lexical_graph.indexing.extract.stages.local_proposition_stage.PropositionExtractor') + def test_as_transform_creates_extractor(self, mock_cls): + """Verify as_transform creates PropositionExtractor.""" + stage = LocalPropositionStage(model_name='test-model', device='cpu') + stage.as_transform() + mock_cls.assert_called_once_with(proposition_model_name='test-model', device='cpu') + + @patch('graphrag_toolkit.lexical_graph.indexing.extract.stages.local_proposition_stage.PropositionExtractor') + def test_as_transform_default_params(self, mock_cls): + """Verify as_transform with no params passes no kwargs.""" + stage = LocalPropositionStage() + stage.as_transform() + mock_cls.assert_called_once_with() + + +class TestLLMTopicExtractionStage: + """Tests for LLMTopicExtractionStage.""" + + def test_is_extraction_stage(self): + """Verify LLMTopicExtractionStage is an ExtractionStage.""" + stage = LLMTopicExtractionStage() + assert isinstance(stage, ExtractionStage) + + def test_input_keys_with_propositions(self): + """Verify input_keys includes PROPOSITIONS_KEY when use_propositions=True.""" + stage = LLMTopicExtractionStage(use_propositions=True) + assert stage.input_keys() == [PROPOSITIONS_KEY] + + def test_input_keys_without_propositions(self): + """Verify input_keys is empty when use_propositions=False.""" + stage = LLMTopicExtractionStage(use_propositions=False) + assert stage.input_keys() == [] + + def test_output_keys(self): + """Verify output_keys returns TOPICS_KEY.""" + stage = LLMTopicExtractionStage() + assert stage.output_keys() == [TOPICS_KEY] + + def test_stage_type(self): + """Verify stage_type is 'llm'.""" + stage = LLMTopicExtractionStage() + assert stage.stage_type == 'llm' + + @patch('graphrag_toolkit.lexical_graph.indexing.extract.stages.llm_topic_extraction_stage.TopicExtractor') + def test_as_transform_with_propositions(self, mock_cls): + """Verify as_transform passes correct source_metadata_field.""" + stage = LLMTopicExtractionStage(use_propositions=True) + stage.as_transform() + mock_cls.assert_called_once_with( + source_metadata_field=PROPOSITIONS_KEY, + prompt_template=None, + llm=None, + entity_classification_provider=None, + topic_provider=None, + schema_constraints='', + ) + + @patch('graphrag_toolkit.lexical_graph.indexing.extract.stages.llm_topic_extraction_stage.TopicExtractor') + def test_as_transform_without_propositions(self, mock_cls): + """Verify as_transform passes None for source_metadata_field.""" + stage = LLMTopicExtractionStage(use_propositions=False) + stage.as_transform() + mock_cls.assert_called_once_with( + source_metadata_field=None, + prompt_template=None, + llm=None, + entity_classification_provider=None, + topic_provider=None, + schema_constraints='', + ) diff --git a/lexical-graph/tests/unit/indexing/test_prompts.py b/lexical-graph/tests/unit/indexing/test_prompts.py index 4872f35fe..42bad25c7 100644 --- a/lexical-graph/tests/unit/indexing/test_prompts.py +++ b/lexical-graph/tests/unit/indexing/test_prompts.py @@ -131,19 +131,23 @@ def test_extract_topics_substitution(self): text = "GraphRAG is a framework." preferred_topics = "Technology, Software" preferred_entity_classifications = "Company, Product" + schema_constraints = "Focus on Technology entities." formatted_prompt = EXTRACT_TOPICS_PROMPT.format( text=text, preferred_topics=preferred_topics, - preferred_entity_classifications=preferred_entity_classifications + preferred_entity_classifications=preferred_entity_classifications, + schema_constraints=schema_constraints ) assert text in formatted_prompt assert preferred_topics in formatted_prompt assert preferred_entity_classifications in formatted_prompt + assert schema_constraints in formatted_prompt assert '{text}' not in formatted_prompt assert '{preferred_topics}' not in formatted_prompt assert '{preferred_entity_classifications}' not in formatted_prompt + assert '{schema_constraints}' not in formatted_prompt def test_domain_entity_classifications_substitution(self): """Verify DOMAIN_ENTITY_CLASSIFICATIONS_PROMPT variable substitution works correctly.""" @@ -250,7 +254,7 @@ def test_prompts_have_no_unintended_placeholders(self): # Define expected placeholders for each prompt expected_placeholders = { 'EXTRACT_PROPOSITIONS_PROMPT': ['{source_info}', '{text}'], - 'EXTRACT_TOPICS_PROMPT': ['{text}', '{preferred_topics}', '{preferred_entity_classifications}'], + 'EXTRACT_TOPICS_PROMPT': ['{text}', '{preferred_topics}', '{preferred_entity_classifications}', '{schema_constraints}'], 'DOMAIN_ENTITY_CLASSIFICATIONS_PROMPT': ['{text_chunks}', '{existing_classifications}'], 'RANK_ENTITY_CLASSIFICATIONS_PROMPT': ['{classifications}'] } diff --git a/lexical-graph/tests/unit/test_lexical_graph_index.py b/lexical-graph/tests/unit/test_lexical_graph_index.py index cd2056d99..95a8d823d 100644 --- a/lexical-graph/tests/unit/test_lexical_graph_index.py +++ b/lexical-graph/tests/unit/test_lexical_graph_index.py @@ -12,6 +12,10 @@ from graphrag_toolkit.lexical_graph import ExtractionConfig from graphrag_toolkit.lexical_graph.lexical_graph_index import LexicalGraphIndex from graphrag_toolkit.lexical_graph.utils.llm_cache import LLMCache +from graphrag_toolkit.lexical_graph.indexing.extract import ( + ExtractionStage, ExtractionSchema, EntityTypeConfig, + LLMPropositionStage, LLMTopicExtractionStage +) class TestExtractionConfig: @@ -46,6 +50,127 @@ def test_extraction_lmm_configured_with_llm_cache_returns_llm_cache(self): assert extraction_config.extraction_llm == llm_cache + def test_default_config_has_no_stages(self): + config = ExtractionConfig() + assert config.stages is None + assert config.schema is None + + def test_from_stages_creates_config_with_stages(self): + stages = [LLMPropositionStage(), LLMTopicExtractionStage()] + config = ExtractionConfig.from_stages(stages=stages) + assert config.stages == stages + assert config.schema is None + assert config.extraction_llm is None + + def test_from_stages_with_schema(self): + stages = [LLMPropositionStage(), LLMTopicExtractionStage()] + schema = ExtractionSchema( + entity_types={'Person': EntityTypeConfig(description='A person')}, + strict=True, + ) + config = ExtractionConfig.from_stages(stages=stages, schema=schema) + assert config.stages == stages + assert config.schema == schema + assert config.schema.strict is True + +class TestCustomPipelineIntegration: + + def test_from_stages_produces_chunking_plus_stage_transforms(self): + """Verify custom stages path includes chunking and stage transforms.""" + from graphrag_toolkit.lexical_graph.indexing.extract import SchemaFilterStage + from graphrag_toolkit.lexical_graph.lexical_graph_index import IndexingConfig + from llama_index.core.node_parser import SentenceSplitter + + schema = ExtractionSchema( + entity_types={'Person': EntityTypeConfig(description='A person')}, + strict=True, + ) + extraction_config = ExtractionConfig.from_stages( + stages=[LLMPropositionStage(), LLMTopicExtractionStage(), SchemaFilterStage(schema=schema)], + schema=schema, + ) + indexing_config = IndexingConfig(extraction=extraction_config) + + graph_store = Mock() + vector_store = Mock() + + with ( + patch( + 'graphrag_toolkit.lexical_graph.lexical_graph_index.GraphStoreFactory.for_graph_store', + return_value=graph_store, + ), + patch( + 'graphrag_toolkit.lexical_graph.lexical_graph_index.MultiTenantGraphStore.wrap', + return_value=graph_store, + ), + patch( + 'graphrag_toolkit.lexical_graph.lexical_graph_index.VectorStoreFactory.for_vector_store', + return_value=vector_store, + ), + patch( + 'graphrag_toolkit.lexical_graph.lexical_graph_index.MultiTenantVectorStore.wrap', + return_value=vector_store, + ), + patch.object(LexicalGraphIndex, '_configure_extraction_pipeline', return_value=([], [])), + ): + index = LexicalGraphIndex(graph_store, vector_store) + + # Call the real method directly + (pre_processors, components) = LexicalGraphIndex._configure_extraction_pipeline(index, indexing_config) + + # Chunking should be first (default SentenceSplitter) + assert isinstance(components[0], SentenceSplitter) + # 3 stage transforms after chunking + assert len(components) == 4 + assert len(pre_processors) == 0 + + def test_from_stages_schema_injected_into_topic_stage(self): + """Verify schema is auto-injected into LLMTopicExtractionStage.""" + from graphrag_toolkit.lexical_graph.lexical_graph_index import IndexingConfig + + schema = ExtractionSchema( + entity_types={'Person': EntityTypeConfig(description='A human being')}, + ) + topic_stage = LLMTopicExtractionStage() + assert topic_stage._schema is None + + extraction_config = ExtractionConfig.from_stages( + stages=[LLMPropositionStage(), topic_stage], + schema=schema, + ) + indexing_config = IndexingConfig(extraction=extraction_config) + + graph_store = Mock() + vector_store = Mock() + + with ( + patch( + 'graphrag_toolkit.lexical_graph.lexical_graph_index.GraphStoreFactory.for_graph_store', + return_value=graph_store, + ), + patch( + 'graphrag_toolkit.lexical_graph.lexical_graph_index.MultiTenantGraphStore.wrap', + return_value=graph_store, + ), + patch( + 'graphrag_toolkit.lexical_graph.lexical_graph_index.VectorStoreFactory.for_vector_store', + return_value=vector_store, + ), + patch( + 'graphrag_toolkit.lexical_graph.lexical_graph_index.MultiTenantVectorStore.wrap', + return_value=vector_store, + ), + patch.object(LexicalGraphIndex, '_configure_extraction_pipeline', return_value=([], [])), + ): + index = LexicalGraphIndex(graph_store, vector_store) + + # Call the real method directly + LexicalGraphIndex._configure_extraction_pipeline(index, indexing_config) + + # Schema should be injected into the topic stage + assert topic_stage._schema == schema + + class TestLexicalGraphIndex: def test_init_invokes_graph_store_init_hook(self):