# Handling Create vs Update Generation Types in Screenshot-to-Code

> Learn how the generationType parameter in screenshot-to-code handles create vs update scenarios. Dynamically adjust code generation, variant counts, and prompts for efficient updates.

- Repository: [Abi Raja/screenshot-to-code](https://github.com/abi/screenshot-to-code)
- Tags: internals
- Published: 2026-03-02

---

**The `generationType` parameter in the WebSocket request determines whether the pipeline performs a fresh code generation or applies targeted updates to existing files, dynamically adjusting variant counts, model selection, and prompt templates.**

The `abi/screenshot-to-code` repository processes UI generation requests through a WebSocket pipeline that branches based on a single `generationType` field. This parameter—accepting either `"create"` or `"update"`—drives the entire execution path from parameter validation through post-processing. Understanding this branching logic is essential for implementing iterative code editing and optimizing generation costs.

## Parameter Extraction and Validation

The flow begins in [`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py) where the system extracts and validates the operation mode.

- **Default behavior**: If the `generationType` field is omitted from the WebSocket request, the system defaults to `"create"` (lines 88-94).
- **Validation**: The pipeline strictly accepts only `"create"` or `"update"` as valid values, rejecting malformed requests early in the parameter extraction stage.

This initial determination persists throughout the middleware chain, influencing downstream decisions without requiring additional user input.

## Variant Count and Status Broadcasting

The generation type immediately affects how many model variants the pipeline executes concurrently:

| Generation Type | Variant Count | Rationale |
|-----------------|---------------|-----------|
| **Create** | Default `NUM_VARIANTS` (typically 3) | Maximizes diversity for fresh designs |
| **Update** | Forced to **2** variants (lines 96-105) | Reduces latency and API costs for iterative edits |

When `generation_type` equals `"update"`, the `StatusBroadcastMiddleware` overrides the default configuration to limit parallel executions, ensuring rapid turnaround for code修改 operations.

## Model Selection Strategy

The `ModelSelectionStage.select_models` method (lines 69-76) implements distinct logic branches based on the operation type.

### Create Flow Selection

For creation tasks, the stage selects `NUM_VARIANTS` models based solely on available API keys (OpenAI, Anthropic, Gemini). The `_get_variant_models` function evaluates key availability to determine the optimal model mix (e.g., all-keys mode, Gemini-only mode).

### Update Flow Optimization

When handling updates with both Gemini and OpenAI keys present, `_get_variant_models` (lines 15-20) returns a specific two-model configuration:

```python
[GEMINI_3_FLASH_PREVIEW_MINIMAL, GPT_5_2_CODEX_LOW]

```

This **Gemini-OpenAI pairing** provides complementary editing capabilities—Gemini handles architectural changes while GPT manages precise code deltas—optimizing for speed and accuracy in modification workflows.

## Prompt Engineering and Context Handling

The `PromptCreationStage` (lines 66-68) tailors instructions based on the `generation_type` parameter, fundamentally changing how the LLM interprets the request.

### Create Mode Prompts

- **Instruction template**: *"Create new UI from this screenshot"*
- **Context**: No previous file state is included
- **Output interpretation**: The LLM generates a complete, standalone HTML/JS bundle

### Update Mode Prompts

- **Instruction template**: *"Update existing code"*
- **Context**: Includes the previous `file_state` object containing current file paths and content
- **Output interpretation**: The LLM produces a **delta** (code diff) that the frontend applies to the existing project structure

This distinction ensures that update operations preserve existing project architecture while modifying only specified components.

## Implementation Examples

### WebSocket Request Payloads

The following JSON structures demonstrate the interface differences between creation and update operations:

```json
// CREATE operation - fresh generation
{
  "generationType": "create",
  "inputMode": "text",
  "prompt": { "text": "Design a landing page for a coffee shop." },
  "generatedCodeConfig": "NEXTJS"
}

```

```json
// UPDATE operation - iterative editing
{
  "generationType": "update",
  "inputMode": "text",
  "prompt": { "text": "Add a sticky header with navigation links." },
  "generatedCodeConfig": "NEXTJS",
  "fileState": {
    "path": "src/pages/index.html",
    "content": "<!DOCTYPE html>..."
  }
}

```

### Pipeline Execution Flow

The orchestration logic in [`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py) processes these requests through distinct stages:

```python

# 1. Extract and validate parameters including generation_type

extracted = await ParameterExtractionStage(...).extract_and_validate(params)

# 2. Broadcast status with variant count adjusted for operation type

await StatusBroadcastMiddleware(...).process(context, next)

# 3. Select models (update -> 2 variants, specific Gemini+OpenAI mix)

variant_models = await ModelSelectionStage(...).select_models(
    generation_type=extracted.generation_type,
    input_mode=extracted.input_mode,
    openai_api_key=extracted.openai_api_key,
    gemini_api_key=GEMINI_API_KEY,
)

# 4. Build targeted prompts based on create vs update logic

prompt_messages = await PromptCreationStage(...).build_prompt_messages(extracted)

# 5. Execute agentic generation across selected variants

variant_completions = await AgenticGenerationStage(...).process_variants(
    variant_models, prompt_messages
)

# 6. Post-process completions (lines 99-104)

await PostProcessingStage().process_completions(
    list(variant_completions.values()), prompt_messages, websocket
)

```

## Key Source Files

Understanding the complete flow requires examining these specific components:

- **[`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py)**: Central orchestration containing all generation-type conditional branches (lines 15-20, 66-76, 88-105).
- **[`backend/prompts/prompt_types.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/prompts/prompt_types.py)**: Defines the `PromptParameters` dataclass including the `generation_type` field.
- **[`backend/prompts/plan.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/prompts/plan.py)**: Implements the planning logic that switches between creation and update strategies.
- **[`backend/prompts/pipeline.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/prompts/pipeline.py)**: Aggregates prompt sections and propagates the `generation_type` flag through the assembly process.
- **[`backend/tests/test_model_selection.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/tests/test_model_selection.py)**: Unit tests verifying correct model list selection for each generation type.
- **[`backend/tests/test_prompts.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/tests/test_prompts.py)**: Validates prompt assembly and `generation_type` flag propagation.

## Summary

- The `generationType` parameter acts as the primary routing mechanism for the entire Screenshot-to-Code pipeline.
- **Create** operations use 3 variants for maximum design diversity, while **Update** operations use 2 variants optimized for speed and cost.
- Update mode employs a specific Gemini-OpenAI model pairing (`GEMINI_3_FLASH_PREVIEW_MINIMAL` and `GPT_5_2_CODEX_LOW`) to generate complementary code deltas.
- Prompt templates dynamically switch between *"Create new UI"* and *"Update existing code"* instructions based on the operation type.
- Update requests require a `file_state` object containing the current code to be modified, which the LLM uses as context for generating precise deltas.

## Frequently Asked Questions

### What happens if I omit the generationType field in my request?

The system defaults to `"create"` mode as implemented in the parameter extraction logic at [`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py) (lines 88-94). This ensures backward compatibility for existing integrations while requiring explicit opt-in for update functionality.

### Why does the update flow restrict variants to 2 instead of 3?

Update operations force exactly 2 variants (overriding the default `NUM_VARIANTS`) to minimize latency and reduce API costs for iterative editing workflows. According to the source code at lines 96-105, this limitation recognizes that modification tasks require less exploratory diversity than initial creation while still providing a fallback option if one model fails.

### How does the file_state parameter function in update mode?

The `file_state` object supplies the current project context to the `PromptCreationStage`, allowing the LLM to generate targeted code deltas rather than complete rewrites. When `generation_type="update"`, the prompt template incorporates this state (lines 66-68), and the `AgenticGenerationStage` interprets the LLM output as modifications to be applied to the existing codebase rather than a fresh bundle.

### Can I customize which models are selected for each generation type?

The current implementation in `_get_variant_models` (lines 15-20) hardcodes the `[GEMINI_3_FLASH_PREVIEW_MINIMAL, GPT_5_2_CODEX_LOW]` pairing for update operations when both API keys are available. For creation tasks, model selection depends on available API keys and the configured mode. Customization would require modifying the model selection logic in [`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py) or extending the `ModelSelectionStage` class.