Per-component prompt configuration

Per-Component Prompt Configuration

Overview

Enable per-component prompt customization within data_types.components[] configuration, allowing each component option to have its own prompt overrides. Additionally, lift the restriction preventing HBCs from being mixed with dynamic components, and validate that HBCs have descriptions when multiple components are configured.

Architecture Changes

1. Configuration Schema Updates

**File: [libs/next_gen_ui_agent/types.py](libs/next_gen_ui_agent/types.py)**

Add a prompt field to AgentConfigComponent (line 68) with proper Field() for JSON schema generation:

class AgentConfigComponent(BaseModel):
    component: str = Field(...)  # existing
    configuration: Optional[AgentConfigDynamicComponentConfiguration] = Field(...)  # existing
    llm_configure: Optional[bool] = Field(...)  # existing
    
    prompt: Optional[AgentConfigPromptComponent] = Field(
        default=None,
        description="Optional prompt customization for this component. Overrides global prompt.components configuration for this component in this data_type context. Has the same fields as AgentConfigPromptComponent. For HBCs in multi-component scenarios, at least 'description' field is required.",
    )
    """Optional prompt customization for this component."""

Important: Use Pydantic's Field() with description parameter to ensure proper JSON schema generation for config validation and IDE support.

This allows configuration like:

data_types:
  movie-list:
    components:
   - component: table
        prompt:
          description: "Custom table description for movies"
   - component: set-of-cards
        prompt:
          description: "Custom cards description"

2. Validation Logic

**File: [libs/next_gen_ui_agent/component_selection_pertype.py](libs/next_gen_ui_agent/component_selection_pertype.py)**

In init_pertype_components_mapping() function (around line 35-92):

• Remove HBC mixing restriction (lines 83-91): Delete the validation that prevents HBCs from being mixed with dynamic components
• Add multi-HBC validation: When multiple components are configured for a data_type, check if any are HBCs. For each HBC, validate that component.prompt.description is defined
• Keep existing validations for llm_configure and configuration fields

Example validation logic:

# Count HBCs in the list
hbc_components = [c for c in components if c.component not in DYNAMIC_COMPONENT_NAMES]

# If multiple components with HBCs, validate descriptions
if len(data_type_config.components) > 1 and hbc_components:
    for hbc in hbc_components:
        if not (hbc.prompt and hbc.prompt.description):
            raise ValueError(
                f"HBC '{hbc.component}' for data type '{data_type}' must have "
                f"prompt.description defined when multiple components are configured"
            )

Note: chart_* and twostep_step2_* fields are not used for HBC prompts - document but don't validate (validation would be complex and unnecessary).

3. Prompt Construction and Caching

**File: [libs/next_gen_ui_agent/component_selection_pertype.py](libs/next_gen_ui_agent/component_selection_pertype.py)**

Major simplification: Delete select_component_with_llm_async() function entirely!

• Move the llm_configure=False merging logic (currently lines 261-278) into the strategy classes
• Strategy now handles this internally in select_component() method after perform_inference() returns
• Delete select_component_with_llm_async() function - no longer needed!

**File: [libs/next_gen_ui_agent/component_selection_llm_onestep.py](libs/next_gen_ui_agent/component_selection_llm_onestep.py)** and **[component_selection_llm_twostep.py](libs/next_gen_ui_agent/component_selection_llm_twostep.py)**

• Store full config: In __init__, store self.config = config (currently only stores selectable_components)
• Add optional data_type parameter to select_component() method (public API)
• Pass data_type down to perform_inference()
• Remove BOTH allowed_components and components_config parameters - no longer needed!
• Update _get_or_build_system_prompt() to accept data_type parameter
• Determine allowed components internally: - When data_type is provided: Extract from self.config.data_types[data_type].components - When data_type is None: Use self.config.selectable_components
• Update cache structure and key logic: - Change cache type hint from dict[frozenset[str], str] to dict[str | None, str] - When data_type is provided: Use data_type string as cache key (e.g., "movies", "products") - When data_type is None (global selection): Use None as cache key (valid in Python dicts)
• Extract prompt overrides from components' .prompt field (when data_type provided)
• Build merged metadata and use it via set_active_component_metadata() before building prompt
• Handle llm_configure=False merging internally in select_component() after inference (moved from select_component_with_llm_async)

Caching strategy rationale:

• Each data_type has exactly ONE set of components configured, so data_type alone is sufficient as cache key
• Global selection (no data_type) has ONE set of selectable components, so None is sufficient
• Much simpler than including components in the key
• Memory overhead is minimal: typically 5-10 data_types + 1 global = ~11 cache entries

4. Metadata Merging

**File: [libs/next_gen_ui_agent/component_metadata.py](libs/next_gen_ui_agent/component_metadata.py)**

Create new function merge_per_component_prompt_overrides():

def merge_per_component_prompt_overrides(
    base_metadata: dict[str, dict[str, Any]],
    components_list: list[AgentConfigComponent]
) -> dict[str, dict[str, Any]]:
    """Merge per-component prompt overrides into metadata.
    
    Args:
        base_metadata: Base metadata (already includes global overrides)
        components_list: List of components with potential prompt overrides
        
    Returns:
        Metadata with per-component overrides applied
    """

This function:

1. Takes base metadata (already has global config.prompt.components applied)
2. Iterates through components_list
3. For each component with .prompt defined, merges those overrides
4. Returns final merged metadata

5. Component Selection Flow (Unified!)

In agent.py select_component() method:

# Try single-component or HBC selection first (no LLM needed)
component = select_component_per_type(input_data, json_data)
if component:
    # ... set metadata and return

# LLM-based selection (unified for both data_type-specific and global)
data_type = input_data.get("type")  
component = await self._component_selection_strategy.select_component(
    inference, 
    user_prompt, 
    input_data_for_strategy,
    data_type=data_type,  # Pass data_type (or None for global)
)

In strategy select_component() method:

1. Determine allowed components based on data_type:

• If data_type: Look up self.config.data_types[data_type].components - If None: Use self.config.selectable_components

2. Extract prompt overrides from components' .prompt field (if data_type)
3. Merge metadata:

• Base COMPONENT_METADATA - Global self.config.prompt.components overrides - Per-component component.prompt overrides (if data_type)

4. Call perform_inference() with data_type for caching
5. Build system prompt using merged metadata and cache it with data_type as key
6. If result has no fields and llm_configure=False, merge with pre-configuration
7. Return complete UIComponentMetadata

Simplification benefits:

• **select_component_with_llm_async() function deleted** - no longer needed!
• No branching in agent.py - single unified path for LLM selection
• Strategy interface is simpler: just data_type parameter (net change: -2 parameters, +1 parameter)
• All data_type and component filtering logic centralized in strategy
• Cleaner, more maintainable code

Testing Strategy

Create comprehensive unit tests in new file: libs/next_gen_ui_agent/component_selection_pertype_test.py

Test cases:

1. Validation tests (test_init_pertype_validation_*):

• Multi-component with HBC without description → ValueError - Multi-component with HBC with description → Success - Single HBC without description → Success (no validation needed) - Multiple HBCs with descriptions → Success - HBC mixed with dynamic components → Success (restriction lifted)

2. Prompt override tests (test_per_component_prompt_*):

• Per-component prompt overrides are applied correctly - Per-component overrides take precedence over global - Multiple components with different prompts - Chart and twostep fields in HBC component prompt (should work, not validated)

3. Caching tests (test_system_prompt_caching_*):

**Test: test_cache_hit_for_same_data_type**

                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            - Create strategy with data_type "movies" configured
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            - Call `_get_or_build_system_prompt("movies")` twice
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            - Assert `_build_system_prompt()` was called only once (using spy/mock)
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            - Assert both calls return the exact same string object (using `is` not `==`)
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            - Assert cache dict has exactly one entry with key `"movies"`

**Test: test_different_cache_entries_per_data_type**

                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                            - Configure two data_types: "movies" and "products" with different prompt overrides
                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                    

…(truncated)