Data Generator

Document Processing Orchestrator Data Generator Package.

This module provides various data generator implementations for different types of data processing.

Classes requiring optional extras are loaded lazily via getattr so that the package can be imported without installing every extra: - ImageCaptionDataGenerator, MultiModelImageCaptionDataGenerator: require the 'image' extra - PIITextAnonymizationDataGenerator: require the 'pii' extra

BaseDataGenerator

Bases: ABC

Base class for data generator.

generate(elements, **kwargs) abstractmethod

Generates data for a list of chunks.

Parameters:

Name	Type	Description	Default
elements	Any	The elements to be used for generating data / metadata. ideally formatted as List[Dict].	required
**kwargs	Any	Additional keyword arguments for customization.	{}

Returns:

Name	Type	Description
Any	Any	The generated data, ideally formatted as List[Dict]. Each dictionary within the list are recommended to follows the structure of model 'Element', to ensure consistency and ease of use across Document Processing Orchestrator.

ChunkRelationMetadataDataGenerator()

Bases: BaseDataGenerator

Data generator that assigns chunk relation metadata to elements.

Each element is treated as an atomic single chunk. The generator populates file_id, chunk_id, previous_chunk, next_chunk, parent_chunk, children_chunk, and order fields on each element's metadata.

Initialize the ChunkRelationMetadataDataGenerator.

generate(elements, **kwargs)

Enrich elements with chunk relation metadata.

Enriches each element with chunk-level relationship metadata: file_id, chunk_id, previous_chunk, next_chunk, parent_chunk, children_chunk, and order.

Processing skip if all chunk relation fields are already present on the first element.

The file_id is resolved in order of priority: 1. file_id kwarg 2. file_id on the first element's metadata 3. Generated UUID value

Parameters:

Name	Type	Description	Default
elements	list[dict[str, Any]]	Elements to enrich with chunk relation metadata.	required
**kwargs	Any	Additional keyword arguments.	{}

Kwargs

file_id (str, optional): Explicit file id. If not provided, falls back to element metadata or a generated UUID when omitted.

Returns:

Type	Description
list[dict[str, Any]]	list[dict[str, Any]]: Elements enriched with chunk relation metadata.

IDRegulationGraphDataGenerator()

Bases: BaseDataGenerator

Maps Indonesian regulation chunks to graph nodes and edges for Neo4j indexing.

Each chunk is enriched with nodes and edges in its metadata. No LLM calls are made; the mapping is purely structural, derived from fields already present in the chunk metadata produced by IDRegulationChunker.

The schema follows ONTOLOGY.md: each article is split into a stable Article identity node and one or more versioned ArticleVersion nodes that carry text and sub-structure (clauses).

Explanation (Penjelasan) chunks become Explanation nodes bound via HAS_EXPLANATION to the ArticleVersion they elucidate, or to the Regulation node for the general (Umum) elucidation.

Concepts and obligations (produced by the optional upstream IDRegulationComplianceDataGenerator) are included when present in the chunk metadata, otherwise skipped.

Initialize IDRegulationGraphDataGenerator.

generate(elements, **kwargs)

Enrich regulation chunks with graph nodes and edges.

Parameters:

Name	Type	Description	Default
elements	list[dict[str, Any]]	Regulation chunks produced by IDRegulationChunker.	required
**kwargs	Any	Unused; accepted for interface compatibility.	{}

Returns:

Type	Description
list[dict[str, Any]]	list[dict[str, Any]]: Same chunks with metadata["nodes"] and metadata["edges"] populated. If the document has no regulation_id (e.g. appendix or FAQ chunks from StructuredElementChunker), all chunks are returned unchanged with empty nodes and edges.

LLMDerivedElementDataGenerator(enabled=False, target_structures=None, derivation_kind=None, system_prompt=None, user_prompt=None, model_id=DEFAULT_MODEL_ID, credentials=None, config=None, max_concurrent_requests=DEFAULT_MAX_CONCURRENT_REQUESTS)

Bases: BaseDataGenerator

Data generator that derives a new element from an eligible existing element via an LLM.

Only eligible elements become targets. Each target produces one derived element, inserted directly after it. Derived elements keep the target's metadata and are marked with their target chunk and derivation kind.

Derivation runs only when both document preconditions hold: 1. Required metadata is present on every element (chunk_id, previous_chunk, next_chunk, and order). 2. The input forms one complete chain in list order (previous_chunk / next_chunk match adjacent elements).

Otherwise the input is returned unchanged and no element is considered eligible.

An element is eligible when it passes every check below, in order: 1. It has no derived_from_chunk (it is not itself a derived element). 2. target_structures is None, or its structure is in target_structures. 3. Its text is non-blank after stripping whitespace. 4. No element in the chain already derives from it with the same derivation_kind.

Initialize the LLMDerivedElementDataGenerator.

Values here are stored as defaults for generate(). Validation runs at call time because each call can override the configuration.

Parameters:

Name	Type	Description	Default
enabled	bool	Whether derived-element generation runs by default. Defaults to False.	False
target_structures	list[str] | None	Structures to derive from; None means all. Defaults to None.	None
derivation_kind	str | None	Label on derived elements. Defaults to None.	None
system_prompt	str | None	System prompt for the LLM. Defaults to None.	None
user_prompt	str | None	User prompt template; must include {text}. Defaults to None.	None
model_id	str | ModelId	Model to use. Defaults to DEFAULT_MODEL_ID.	DEFAULT_MODEL_ID
credentials	str | dict[str, Any] | None	Model credentials. Defaults to None.	None
config	dict[str, Any] | None	Extra options forwarded to the LLM invoker. Defaults to None.	None
max_concurrent_requests	int	Max concurrent LLM requests per generate call. Defaults to DEFAULT_MAX_CONCURRENT_REQUESTS.	DEFAULT_MAX_CONCURRENT_REQUESTS

generate(elements, **kwargs)

Generate and insert derived elements after their target elements.

When enabled, each eligible target gets one LLM-derived element inserted directly after it. Bad configuration raises ValueError; request-time failures are logged and that element is skipped.

Parameters:

Name	Type	Description	Default
elements	list[dict[str, Any]]	Elements to process. Each element must have chunk_id, previous_chunk, next_chunk, and order, and the elements must form one complete chain in list order.	required
**kwargs	Any	Optional configuration overrides for this call.	{}

Kwargs

enabled (bool, optional): Whether derived-element generation runs by default. Defaults to __init__ value. target_structures (list[str] | None, optional): Structures to derive from; None means all. Defaults to __init__ value. derivation_kind (str | None, optional): Label on derived elements. Defaults to __init__ value. system_prompt (str | None, optional): System prompt for the LLM. Defaults to __init__ value. user_prompt (str | None, optional): User prompt template; must include {text}. Defaults to __init__ value. model_id (str | ModelId, optional): Model to use. Defaults to __init__ value. credentials (str | dict[str, Any] | None, optional): Model credentials. Defaults to __init__ value. config (dict[str, Any] | None, optional): Extra options forwarded to the LLM invoker. Defaults to __init__ value. max_concurrent_requests (int, optional): Max concurrent LLM requests per generate call. Defaults to __init__ value.

Returns:

Type	Description
list[dict[str, Any]]	list[dict[str, Any]]: Elements with derived items inserted after their targets, or the input unchanged when nothing is derived. The input list is never mutated.

Raises:

Type	Description
ValueError	If enabled and a setting is missing/invalid, or the LLM processor cannot be built.

LLMTextRewriteDataGenerator(default_model_id=DEFAULT_MODEL_ID, default_system_prompt=DEFAULT_SYSTEM_PROMPT, default_user_prompt=DEFAULT_USER_PROMPT, default_text_rewrite_enabled=True, default_structures_to_rewrite=None, default_hyperparameters=None, default_retry_config=None, default_include_media_images_as_context=True, max_concurrent_requests=DEFAULT_MAX_CONCURRENT_REQUESTS)

Bases: BaseDataGenerator

LLM-powered element text rewrite data generator with lazy initialization and batching.

This generator rewrites the text field of document elements using an LLM while preserving element structure and metadata. It supports dynamic model selection, configuration-based processor caching, concurrent batch processing, and optional image context for multimodal models.

Designed primarily for document normalization tasks such as OCR cleanup, formatting correction, and layout refinement.

Initialize the LLMTextRewriteDataGenerator.

Parameters:

Name	Type	Description	Default
default_model_id	str	Default model ID used when model_id is not provided to generate(). Defaults to DEFAULT_MODEL_ID.	DEFAULT_MODEL_ID
default_system_prompt	str	Default system prompt used when system_prompt is not provided to generate(). Defaults to DEFAULT_SYSTEM_PROMPT.	DEFAULT_SYSTEM_PROMPT
default_user_prompt	str	Default user prompt used when user_prompt is not provided to generate(). Defaults to DEFAULT_USER_PROMPT.	DEFAULT_USER_PROMPT
default_text_rewrite_enabled	bool	Default value for text_rewrite_enabled in generate(). Defaults to True.	True
default_structures_to_rewrite	list[str] | None	Default value for structures_to_rewrite in generate(). If None, all structures are eligible for rewriting. Defaults to None.	None
default_hyperparameters	dict[str, Any] | None	Default hyperparameters passed to the LMInvoker when default_hyperparameters is not provided in generate(). Defaults to None.	None
default_retry_config	dict[str, Any] | None	Default retry configuration passed to the LMInvoker when retry_config is not provided in generate(). Defaults to None.	None
default_include_media_images_as_context	bool	Default value for include_media_images_as_context in generate(). Defaults to True.	True
max_concurrent_requests	int	Default maximum number of concurrent LLM requests per generate() call when max_concurrent_requests is not provided to generate(). Defaults to DEFAULT_MAX_CONCURRENT_REQUESTS.	DEFAULT_MAX_CONCURRENT_REQUESTS

Raises:

Type	Description
ValueError	If max_concurrent_requests is less than 1.

generate(elements, **kwargs)

Rewrite element.text for elements matching structures_to_rewrite using an LLM.

Parameters:

Name	Type	Description	Default
elements	list[dict[str, Any]]	List of dictionaries containing elements to be processed.	required
**kwargs	Any	Additional keyword arguments for the LLM text rewrite process.	{}

Kwargs

text_rewrite_enabled (bool, optional): Whether to enable LLM text rewriting. Defaults to True or the value configured in init. structures_to_rewrite (list[str], optional): List of element structure types whose text field should be rewritten. Elements with structures not in this list are passed through unchanged. Defaults to the value configured in init (applies rewriting to all structures if None). model_id (str, optional): The ID of the model to use for text rewriting. Defaults to DEFAULT_MODEL_ID or the value configured in init. system_prompt (str, optional): The system prompt to use for text rewriting. Defaults to DEFAULT_SYSTEM_PROMPT or the value configured in init. user_prompt (str, optional): The user prompt template for text rewriting. Must contain a {text} placeholder. Defaults to DEFAULT_USER_PROMPT or the value configured in init. default_hyperparameters (dict[str, Any], optional): Additional hyperparameters passed to the LMInvoker configuration. Defaults to {} or the value configured in init. retry_config (dict[str, Any], optional): Retry configuration passed to the LMInvoker. Defaults to {} or the value configured in init. include_media_images_as_context (bool, optional): Whether to attach associated media images from element.metadata.media as visual context when invoking the LLM. Defaults to True or the value configured in init. max_concurrent_requests (int, optional): Maximum number of LLM requests in flight at once for this generate() call. Defaults to the value configured in init. api_key (str, optional): API key for the model. Defaults to None, in which case no credentials are passed and the LMInvoker falls back to the provider environment variable.

Returns:

Type	Description
list[dict[str, Any]]	list[dict[str, Any]]: List of dictionaries with text rewritten for matching elements. Non-matching elements are returned unchanged.

Raises:

Type	Description
ValueError	If max_concurrent_requests is less than 1.

MetadataEnrichmentDataGenerator(kwargs_to_unpack_into_metadata=None, kwargs_to_keep_as_metadata_value=None)

Bases: BaseDataGenerator

Data generator that enriches every element with metadata supplied by the caller.

Both are configured at construction time. kwargs_to_unpack_into_metadata names kwargs that each hold a whole metadata dictionary, unpacked as is to form the base. kwargs_to_keep_as_metadata_value names kwargs whose value is kept as one metadata entry, keyed by that same name, on top of that base. Later values win over earlier ones.

Values that are None are ignored. Existing metadata is kept unless a non-None value is provided for the same key, in which case it is overwritten.

Attributes:

Name	Type	Description
kwargs_to_unpack_into_metadata	tuple[str, ...]	Kwarg names each holding a metadata dictionary, unpacked in order to form the base metadata.
kwargs_to_keep_as_metadata_value	tuple[str, ...]	Kwarg names whose value is kept on top of that base as one entry, keyed by that same name.

Initialize the data generator.

Parameters:

Name	Type	Description	Default
kwargs_to_unpack_into_metadata	Sequence[str] | None	Kwarg names each holding a metadata dictionary, unpacked in order to form the base metadata. Defaults to None.	None
kwargs_to_keep_as_metadata_value	Sequence[str] | None	Kwarg names whose value is kept on top of that base as one entry, keyed by that same name. Defaults to None.	None

generate(elements, **kwargs)

Add caller-provided metadata to every element.

Parameters:

Name	Type	Description	Default
elements	list[dict[str, Any]]	The elements to enrich, each following the Element model.	required
**kwargs	Any	Additional keyword arguments, read through the configured keys.	{}

Returns:

Type	Description
list[dict[str, Any]]	list[dict[str, Any]]: The elements with the caller-provided metadata added.

__getattr__(name)

Lazily import classes that depend on optional extras.

Parameters:

Name	Type	Description	Default
name	str	The attribute name being accessed on this module.	required

Returns:

Name	Type	Description
type	type	The requested class, imported from its submodule on first access.

Raises:

Type	Description
ImportError	If the class is found in _OPTIONAL_IMPORTS but its required extra is not installed.
AttributeError	If name is not a known attribute of this module.