hcie-rust-v3.05/hcie-iced-app/crates/hcie-iced-gui/src/ai_script.rs

//! AI Script Generator — LLM-powered procedural drawing script generation.
//!
//! Provides async functions to call Ollama, Claude, and OpenAI APIs to generate
//! HCIE-Script DSL commands from natural language descriptions. Also includes
/// the system prompt, provider configuration, and output cleaning utilities.

/// System prompt sent to LLMs when generating HCIE-Script DSL commands.
pub const SYSTEM_PROMPT: &str = r#"You are HCIE-Script, a code generator for a procedural drawing DSL.
Output ONLY valid HCIE-Script DSL commands with no explanation, no markdown, no code fences.

Available commands:
- layer Name — create raster layer
- select_layer Name
- brush STYLE — round, square, pencil, inkpen, charcoal, watercolor, marker, oil, airbrush, spray, clouds, dirt, tree, meadow, rock, bristle, leaf, wetpaint, sketch, hatch, calligraphy, blender, mixer, glow, crayon
- color #RRGGBB — set foreground color
- size N — brush size in pixels
- opacity N — 0..1
- hardness N — 0..1
- flow N — 0..2
- pressure N / pressure_start N / pressure_end N — 0..1
- wiggle N — per-point jitter
- draw x1,y1 x2,y2 ... — multi-point stroke (coords in pixels)
- line x1,y1 x2,y2
- rect x1,y1 x2,y2 — filled rectangle in CURRENT color
- circle cx,cy — single brush dab
- triangle x1,y1 x2,y2 x3,y3 — vector triangle
- bezier x1,y1 cx1,cy1 cx2,cy2 x2,y2 [steps]
- scatter N x,y w,h [size] — golden-angle scatter
- gradient x1,y1 x2,y2 #hex1 #hex2 ... N
- set key=value key=value — multi-set
- var name value — declare variables, reference as $name
- repeat N { ... } — loop with $i (index), $n (count)
- blend MODE — normal, multiply, overlay, screen, etc.

Golden rules:
1. Create a layer BEFORE drawing on it.
2. Use brush BEFORE draw/line/bezier.
3. Drawing commands operate on the CURRENT layer and CURRENT brush state.
4. Coordinates are pixel values (0..800 width, 0..600 height).
5. Use gradient for smooth skies and backgrounds.
6. Use scatter for grass, flowers, stars, foliage.
7. Use bezier/curve for smooth natural lines.
8. Build scenes back-to-front: sky → background → midground → foreground.
"#;

/// LLM provider configuration.
pub struct Provider {
    pub base_url: String,
    pub model: String,
}

/// Generate a script via the Ollama /api/chat endpoint (async).
///
/// Sends the system prompt and user prompt as a chat completion request with
/// `stream: false`. Returns the cleaned DSL script text on success.
pub async fn generate_via_ollama(provider: &Provider, prompt: &str) -> Result<String, String> {
    let url = format!("{}/api/chat", provider.base_url.trim_end_matches('/'));
    let body = serde_json::json!({
        "model": provider.model,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt}
        ],
        "stream": false,
        "options": { "num_predict": 2048 }
    });

    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(120))
        .build()
        .map_err(|e| format!("Client error: {}", e))?;

    let resp = client
        .post(&url)
        .json(&body)
        .send()
        .await
        .map_err(|e| format!("Request failed: {}", e))?;

    if !resp.status().is_success() {
        return Err(format!("API error: {}", resp.status()));
    }

    let json: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| format!("Parse error: {}", e))?;
    let content = json["message"]["content"]
        .as_str()
        .ok_or_else(|| "No content in response".to_string())?;

    Ok(clean_script_output(content))
}

/// Generate a script via the Claude /v1/messages endpoint (async).
///
/// Uses the Anthropic Messages API format with the system prompt as a
/// top-level `system` field. Returns the cleaned DSL script text.
pub async fn generate_via_claude(provider: &Provider, prompt: &str) -> Result<String, String> {
    let url = format!("{}/v1/messages", provider.base_url.trim_end_matches('/'));
    let body = serde_json::json!({
        "model": provider.model,
        "max_tokens": 2048,
        "system": SYSTEM_PROMPT,
        "messages": [
            {"role": "user", "content": prompt}
        ]
    });

    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(120))
        .build()
        .map_err(|e| format!("Client error: {}", e))?;

    let api_key = std::env::var("ANTHROPIC_API_KEY")
        .map_err(|_| "ANTHROPIC_API_KEY is not set".to_string())?;
    let resp = client
        .post(&url)
        .header("x-api-key", api_key)
        .header("anthropic-version", "2023-06-01")
        .json(&body)
        .send()
        .await
        .map_err(|e| format!("Request failed: {}", e))?;

    if !resp.status().is_success() {
        let status = resp.status();
        let text = resp.text().await.unwrap_or_default();
        return Err(format!("API error {}: {}", status, text));
    }

    let json: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| format!("Parse error: {}", e))?;
    let content = json["content"][0]["text"]
        .as_str()
        .ok_or_else(|| "No content in response".to_string())?;

    Ok(clean_script_output(content))
}

/// Generate a script via the OpenAI /v1/chat/completions endpoint (async).
///
/// Uses the standard Chat Completions format with the system prompt as the
/// first message. Returns the cleaned DSL script text.
pub async fn generate_via_openai(provider: &Provider, prompt: &str) -> Result<String, String> {
    let url = format!(
        "{}/chat/completions",
        provider.base_url.trim_end_matches('/')
    );
    let body = serde_json::json!({
        "model": provider.model,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt}
        ],
        "max_tokens": 2048
    });

    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(120))
        .build()
        .map_err(|e| format!("Client error: {}", e))?;

    let api_key =
        std::env::var("OPENAI_API_KEY").map_err(|_| "OPENAI_API_KEY is not set".to_string())?;
    let resp = client
        .post(&url)
        .bearer_auth(api_key)
        .json(&body)
        .send()
        .await
        .map_err(|e| format!("Request failed: {}", e))?;

    if !resp.status().is_success() {
        let status = resp.status();
        let text = resp.text().await.unwrap_or_default();
        return Err(format!("API error {}: {}", status, text));
    }

    let json: serde_json::Value = resp
        .json()
        .await
        .map_err(|e| format!("Parse error: {}", e))?;
    let content = json["choices"][0]["message"]["content"]
        .as_str()
        .ok_or_else(|| "No content in response".to_string())?;

    Ok(clean_script_output(content))
}

/// Strip markdown code fences and leading/trailing whitespace from LLM output.
///
/// LLMs often wrap DSL output in ```dsl, ```script, ```hcie, or plain ``` fences.
/// This function removes those wrappers to extract the raw script text.
pub fn clean_script_output(text: &str) -> String {
    let text = text.trim();
    let text = text.strip_prefix("```dsl").unwrap_or(text);
    let text = text.strip_prefix("```script").unwrap_or(text);
    let text = text.strip_prefix("```hcie").unwrap_or(text);
    let text = text.strip_prefix("```").unwrap_or(text);
    let text = text.strip_suffix("```").unwrap_or(text);
    text.trim().to_string()
}

/// Validate a generated payload before exposing Run as an executable action.
pub fn validate_script_output(text: &str) -> Result<String, String> {
    let cleaned = clean_script_output(text);
    if cleaned.is_empty() {
        return Err("Provider returned an empty script".to_string());
    }
    if cleaned.len() > 512 * 1024 {
        return Err("Generated script exceeds the 512 KiB safety limit".to_string());
    }
    crate::script::parser::parse_dsl_script(&cleaned)
        .map_err(|error| format!("Generated script is invalid: {error}"))?;
    Ok(cleaned)
}

/// Generate a script using the given provider type and configuration (async).
///
/// Dispatches to the appropriate provider function based on the `provider_type`
/// string ("ollama", "claude", or "openai"). Defaults to Ollama.
pub async fn generate_script(
    provider_type: &str,
    provider: &Provider,
    prompt: &str,
) -> Result<String, String> {
    match provider_type {
        "claude" => generate_via_claude(provider, prompt).await,
        "openai" => generate_via_openai(provider, prompt).await,
        "ollama" => generate_via_ollama(provider, prompt).await,
        unsupported => Err(format!("Unsupported AI provider: {unsupported}")),
    }
    .and_then(|output| validate_script_output(&output))
}

#[cfg(test)]
mod tests {
    use super::validate_script_output;

    #[test]
    fn generated_output_must_parse_before_run_is_enabled() {
        assert!(
            validate_script_output("```dsl\nlayer Sky\ncolor #112233\nrect 0,0 10,10\n```").is_ok()
        );
        assert!(validate_script_output("paint magically")
            .unwrap_err()
            .contains("invalid"));
    }
}