Genarrative/server-rs/crates/api-server/src/match3d/vector_engine_gemini.rs

use super::*;

pub(super) fn build_match3d_vector_engine_gemini_image_request_body(
    prompt: &str,
    negative_prompt: &str,
    aspect_ratio: &str,
) -> Value {
    json!({
        "contents": [{
            "role": "user",
            "parts": [{
                "text": build_match3d_vector_engine_gemini_prompt(prompt, negative_prompt),
            }],
        }],
        "generationConfig": {
            "responseModalities": ["TEXT", "IMAGE"],
            "imageConfig": {
                "aspectRatio": aspect_ratio,
            },
        },
    })
}

pub(super) fn build_match3d_vector_engine_gemini_generate_content_url(
    settings: &Match3DVectorEngineGeminiImageSettings,
) -> String {
    let base_url = settings.base_url.trim_end_matches("/v1");
    format!(
        "{}/v1beta/models/{}:generateContent",
        base_url, MATCH3D_MATERIAL_VECTOR_ENGINE_GEMINI_MODEL
    )
}

fn build_match3d_vector_engine_gemini_prompt(prompt: &str, negative_prompt: &str) -> String {
    let prompt = prompt.trim();
    let negative_prompt = negative_prompt.trim();
    if negative_prompt.is_empty() {
        return prompt.to_string();
    }

    format!("{prompt}\n避免：{negative_prompt}")
}

pub(super) fn extract_match3d_b64_images(payload: &Value) -> Vec<String> {
    let mut values = Vec::new();
    collect_match3d_strings_by_key(payload, "b64_json", &mut values);
    collect_match3d_inline_image_data(payload, &mut values);
    values
}

fn collect_match3d_inline_image_data(payload: &Value, results: &mut Vec<String>) {
    match payload {
        Value::Array(entries) => {
            for entry in entries {
                collect_match3d_inline_image_data(entry, results);
            }
        }
        Value::Object(object) => {
            for key in ["inlineData", "inline_data"] {
                if let Some(Value::Object(inline_data)) = object.get(key) {
                    let mime_type = inline_data
                        .get("mimeType")
                        .or_else(|| inline_data.get("mime_type"))
                        .and_then(Value::as_str)
                        .map(str::trim)
                        .unwrap_or("image/png")
                        .to_ascii_lowercase();
                    if !mime_type.is_empty() && !mime_type.starts_with("image/") {
                        continue;
                    }
                    if let Some(data) = inline_data
                        .get("data")
                        .and_then(Value::as_str)
                        .map(str::trim)
                        .filter(|value| !value.is_empty())
                    {
                        results.push(data.to_string());
                    }
                }
            }
            for nested_value in object.values() {
                collect_match3d_inline_image_data(nested_value, results);
            }
        }
        _ => {}
    }
}

fn collect_match3d_strings_by_key(payload: &Value, target_key: &str, results: &mut Vec<String>) {
    match payload {
        Value::Array(entries) => {
            for entry in entries {
                collect_match3d_strings_by_key(entry, target_key, results);
            }
        }
        Value::Object(object) => {
            for (key, nested_value) in object {
                if key == target_key {
                    match nested_value {
                        Value::String(text) => {
                            let text = text.trim();
                            if !text.is_empty() {
                                results.push(text.to_string());
                            }
                        }
                        Value::Array(entries) => {
                            for entry in entries {
                                if let Some(text) = entry
                                    .as_str()
                                    .map(str::trim)
                                    .filter(|value| !value.is_empty())
                                {
                                    results.push(text.to_string());
                                }
                            }
                        }
                        _ => {}
                    }
                }
                collect_match3d_strings_by_key(nested_value, target_key, results);
            }
        }
        _ => {}
    }
}