Add Augmentation Strategy

Step 1: Extraction          Step 2: Augmentation
    Text                        Initial Triples
      │                              │
      ▼                              ▼
extract_triples()           augment_triples() orchestrator
      │                       ┌──────┴──────┐
      ▼                       ▼             ▼
Initial Triples        connectivity    your_strategy
                              │             │
                              └──────┬──────┘
                                     ▼
                              Refined Triples + Metadata

class AugmentationStrategy(Protocol):
    def __call__(
        self,
        client: BaseLLMClient,
        domain: KnowledgeDomain,
        text: str,
        triples: list[Triple],
        **kwargs: Any
    ) -> tuple[list[Triple], dict[str, Any]]:
        """Returns (refined_triples, metadata)."""
        ...

@register_strategy("enrichment")
def enrichment_strategy(
    client: BaseLLMClient,
    domain: KnowledgeDomain,
    text: str,
    triples: list[Triple],
    *,
    max_iterations: int = 3,
    temperature: float = 0.0,
    max_tokens: int | None = None,
    augmentation_prompt_override: str | None = None,
    **kwargs: Any
) -> tuple[list[Triple], dict[str, Any]]:
    """Enrichment augmentation: Add missing entity attributes and relations.

    Args:
        client: LLM client for generation
        domain: Knowledge domain with prompts/examples
        text: Source text to analyze
        triples: Initial triples from extraction
        max_iterations: Max refinement iterations (default: 3)
        temperature: Sampling temperature for LLM (0.0 = deterministic)
        max_tokens: Max tokens for LLM
        augmentation_prompt_override: Override the default prompt

    Returns:
        Tuple of (enriched_triples, metadata)
    """
    # 1. Fetch strategy-specific resources from domain
    augmentation_component = domain.get_augmentation("enrichment")
    aug_prompt_template = augmentation_prompt_override or augmentation_component.prompt
    constraints = collect_schema_constraints(domain, augmentation_component.examples)

    all_triples = list(triples)  # Copy to avoid mutation
    iterations_data = []
    error_occurred = False

    # 2. Iteration loop with error preservation
    for i in range(max_iterations):
        try:
            # Build prompt with current state
            current_triples_dicts = [t.model_dump() for t in all_triples]
            record = {
                "text": text,
                "current_triples": current_triples_dicts,
            }

            final_prompt = render_prompt_template(
                aug_prompt_template,
                record,
                schema_guidance=build_schema_guidance(constraints),
            )

            # Call client.augment() — NOT extract()
            new_triples_raw = client.augment(
                text=final_prompt,
                prompt_description="Enrich entities with missing attributes and relations",
                format_type=Triple,
                temperature=temperature,
                max_tokens=max_tokens
            )

            # 3. Validate and force CONTEXTUAL inference type
            new_triples = []
            normalized_raw: list[dict[str, Any]] = []
            for t_raw in new_triples_raw:
                try:
                    t_dict = t_raw if isinstance(t_raw, dict) else t_raw.model_dump()
                    t_dict["inference"] = InferenceType.CONTEXTUAL
                    new_triples.append(Triple(**t_dict))
                    normalized_raw.append(t_dict)
                except Exception as e:
                    print(f"Warning: Skipping invalid augmented triple: {e}")
                    continue

            # 4. Schema validation
            validated_triples, validation_summary = validate_triples_against_schema(
                new_triples,
                constraints,
                raw_triples=normalized_raw,
            )
            warn_on_schema_validation("enrichment", validation_summary)

            all_triples.extend(validated_triples)
            iterations_data.append({
                "iteration": i + 1,
                "new_triples_count": len(validated_triples),
                "status": "success",
                "schema_validation": validation_summary,
            })

        except Exception as e:
            print(f"Error during enrichment iteration {i+1}: {e}")
            iterations_data.append({
                "iteration": i + 1,
                "status": "failed",
                "error": str(e)
            })
            error_occurred = True
            break

    metadata = {
        "strategy": "enrichment",
        "iterations": iterations_data,
        "partial_result": error_occurred,
        "schema_constraints_applied": constraints.enforce,
        "allowed_entity_types": list(constraints.entity_types),
        "allowed_relation_types": list(constraints.relation_types),
    }

    return all_triples, metadata

kgb/domains/<domain>/augmentation/enrichment/
├── prompt.md          # Strategy-specific prompt (must be .md)
└── examples.json      # Few-shot augmentation examples

[
  {
    "input": {"text": "...", "entities": ["..."]},
    "output": [{"head": "...", "relation": "...", "tail": "...", "inference": "contextual", "justification": "..."}]
  }
]