From 3ff24e060e65d1afd88700f743afa31db74e45d1 Mon Sep 17 00:00:00 2001 From: liubu Date: Sat, 3 Jan 2026 23:34:03 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=20Antigravity=20?= =?UTF-8?q?=E5=9B=BE=E5=83=8F=E7=94=9F=E6=88=90=20API=20(OpenAI=20?= =?UTF-8?q?=E5=85=BC=E5=AE=B9)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 实现 /v1/images/generations 端点 - 添加 ImageGenerationRequest/Response 数据模型 - 实现 OpenAI 到 Antigravity 请求/响应转换器 - 支持 dall-e-3/dall-e-2 模型名映射到 gemini-3-pro-image - 支持 b64_json 和 url 两种响应格式 - 添加 Python SDK 集成测试脚本 --- scripts/test_image_api.py | 369 ++++++++++ .../src/converter/openai_to_antigravity.rs | 696 +++++++++++++++++- src-tauri/src/models/openai.rs | 79 ++ src-tauri/src/providers/antigravity.rs | 12 + .../src/server/handlers/image_handler.rs | 352 +++++++++ src-tauri/src/server/handlers/mod.rs | 2 + src-tauri/src/server/mod.rs | 5 + 7 files changed, 1491 insertions(+), 24 deletions(-) create mode 100755 scripts/test_image_api.py create mode 100644 src-tauri/src/server/handlers/image_handler.rs diff --git a/scripts/test_image_api.py b/scripts/test_image_api.py new file mode 100755 index 000000000..dce674c18 --- /dev/null +++ b/scripts/test_image_api.py @@ -0,0 +1,369 @@ +#!/usr/bin/env python3 +""" +OpenAI 兼容图像生成 API 测试脚本 + +使用 OpenAI Python SDK 测试 Antigravity 图像生成 API。 + +使用方法: + # 安装依赖 + pip install openai + + # 运行测试（需要先启动 API Server） + python scripts/test_image_api.py + + # 指定自定义 API 地址和密钥 + python scripts/test_image_api.py --base-url http://localhost:8999 --api-key your-key + +环境变量: + PROXYCAST_BASE_URL: API 服务器地址（默认: http://localhost:8999） + PROXYCAST_API_KEY: API 密钥（默认: pc_LXZbIv3o78WpHuQwqgmwC0U4G0cY5UtQ） +""" + +import argparse +import base64 +import os +import sys +from datetime import datetime + +try: + from openai import OpenAI +except ImportError: + print("错误: 请先安装 openai 库") + print("运行: pip install openai") + sys.exit(1) + + +def test_image_generation_url(client: OpenAI, prompt: str) -> bool: + """ + 测试 URL 响应格式的图像生成 + + Args: + client: OpenAI 客户端 + prompt: 图像生成提示词 + + Returns: + 测试是否通过 + """ + print("\n" + "=" * 60) + print("测试 1: URL 响应格式") + print("=" * 60) + print(f"提示词: {prompt}") + + try: + response = client.images.generate( + model="dall-e-3", # 会被映射到 gemini-3-pro-image + prompt=prompt, + n=1, + size="1024x1024", + response_format="url" + ) + + # 验证响应结构 + print(f"\n响应时间戳: {response.created}") + print(f"生成图片数量: {len(response.data)}") + + if len(response.data) == 0: + print("❌ 错误: 没有生成图片") + return False + + image = response.data[0] + + # 验证 URL 格式 + if image.url: + print(f"URL 长度: {len(image.url)} 字符") + if image.url.startswith("data:image/"): + print("✅ URL 格式正确 (data URL)") + else: + print(f"⚠️ URL 格式: {image.url[:50]}...") + else: + print("❌ 错误: URL 为空") + return False + + # 验证 revised_prompt + if image.revised_prompt: + print(f"修订提示词: {image.revised_prompt[:100]}...") + else: + print("ℹ️ 没有修订提示词") + + print("\n✅ 测试 1 通过") + return True + + except Exception as e: + print(f"\n❌ 测试 1 失败: {e}") + return False + + +def test_image_generation_b64(client: OpenAI, prompt: str) -> bool: + """ + 测试 b64_json 响应格式的图像生成 + + Args: + client: OpenAI 客户端 + prompt: 图像生成提示词 + + Returns: + 测试是否通过 + """ + print("\n" + "=" * 60) + print("测试 2: b64_json 响应格式") + print("=" * 60) + print(f"提示词: {prompt}") + + try: + response = client.images.generate( + model="gemini-3-pro-image-preview", # 直接使用 Gemini 模型名 + prompt=prompt, + n=1, + response_format="b64_json" + ) + + # 验证响应结构 + print(f"\n响应时间戳: {response.created}") + print(f"生成图片数量: {len(response.data)}") + + if len(response.data) == 0: + print("❌ 错误: 没有生成图片") + return False + + image = response.data[0] + + # 验证 b64_json 格式 + if image.b64_json: + print(f"Base64 数据长度: {len(image.b64_json)} 字符") + + # 尝试解码验证 + try: + decoded = base64.b64decode(image.b64_json) + print(f"解码后大小: {len(decoded)} 字节") + + # 检查图片魔数 + if decoded[:8] == b'\x89PNG\r\n\x1a\n': + print("✅ 图片格式: PNG") + elif decoded[:2] == b'\xff\xd8': + print("✅ 图片格式: JPEG") + elif decoded[:4] == b'GIF8': + print("✅ 图片格式: GIF") + elif decoded[:4] == b'RIFF': + print("✅ 图片格式: WebP") + else: + print(f"⚠️ 未知图片格式: {decoded[:8].hex()}") + + except Exception as e: + print(f"⚠️ Base64 解码失败: {e}") + else: + print("❌ 错误: b64_json 为空") + return False + + # 验证 revised_prompt + if image.revised_prompt: + print(f"修订提示词: {image.revised_prompt[:100]}...") + else: + print("ℹ️ 没有修订提示词") + + print("\n✅ 测试 2 通过") + return True + + except Exception as e: + print(f"\n❌ 测试 2 失败: {e}") + return False + + +def test_error_handling(client: OpenAI) -> bool: + """ + 测试错误处理 + + Args: + client: OpenAI 客户端 + + Returns: + 测试是否通过 + """ + print("\n" + "=" * 60) + print("测试 3: 错误处理") + print("=" * 60) + + try: + # 测试空提示词 + print("测试空提示词...") + try: + response = client.images.generate( + model="dall-e-3", + prompt="", # 空提示词 + n=1 + ) + print("❌ 错误: 应该拒绝空提示词") + return False + except Exception as e: + error_msg = str(e).lower() + if "prompt" in error_msg or "empty" in error_msg or "required" in error_msg: + print(f"✅ 正确拒绝空提示词: {e}") + else: + print(f"⚠️ 收到错误但消息不明确: {e}") + + print("\n✅ 测试 3 通过") + return True + + except Exception as e: + print(f"\n❌ 测试 3 失败: {e}") + return False + + +def test_response_structure(client: OpenAI, prompt: str) -> bool: + """ + 测试响应结构符合 OpenAI 规范 + + Args: + client: OpenAI 客户端 + prompt: 图像生成提示词 + + Returns: + 测试是否通过 + """ + print("\n" + "=" * 60) + print("测试 4: 响应结构验证") + print("=" * 60) + print(f"提示词: {prompt}") + + try: + response = client.images.generate( + model="dall-e-3", + prompt=prompt, + n=1, + response_format="url" + ) + + # 验证 created 字段 + if response.created: + print(f"✅ created 字段存在: {response.created}") + # 验证是有效的 Unix 时间戳 + if response.created > 0: + dt = datetime.fromtimestamp(response.created) + print(f" 时间: {dt}") + else: + print("❌ created 不是有效时间戳") + return False + else: + print("❌ created 字段缺失") + return False + + # 验证 data 字段 + if response.data is not None: + print(f"✅ data 字段存在: {len(response.data)} 项") + if len(response.data) > 0: + print("✅ data 数组非空") + else: + print("❌ data 数组为空") + return False + else: + print("❌ data 字段缺失") + return False + + # 验证每个图片项 + for i, image in enumerate(response.data): + print(f"\n图片 {i + 1}:") + has_url = image.url is not None + has_b64 = image.b64_json is not None + + if has_url: + print(f" ✅ url 字段存在") + if has_b64: + print(f" ✅ b64_json 字段存在") + + if not has_url and not has_b64: + print(f" ❌ 缺少 url 和 b64_json") + return False + + if image.revised_prompt: + print(f" ✅ revised_prompt 字段存在") + else: + print(f" ℹ️ revised_prompt 字段为空") + + print("\n✅ 测试 4 通过") + return True + + except Exception as e: + print(f"\n❌ 测试 4 失败: {e}") + return False + + +def main(): + parser = argparse.ArgumentParser( + description="测试 OpenAI 兼容图像生成 API" + ) + parser.add_argument( + "--base-url", + default=os.environ.get("PROXYCAST_BASE_URL", "http://localhost:8999"), + help="API 服务器地址" + ) + parser.add_argument( + "--api-key", + default=os.environ.get("PROXYCAST_API_KEY", "pc_LXZbIv3o78WpHuQwqgmwC0U4G0cY5UtQ"), + help="API 密钥" + ) + parser.add_argument( + "--prompt", + default="A cute fluffy cat sitting on a windowsill, looking at the sunset", + help="测试用的图像生成提示词" + ) + parser.add_argument( + "--skip-generation", + action="store_true", + help="跳过实际图像生成测试（仅测试错误处理）" + ) + + args = parser.parse_args() + + print("=" * 60) + print("OpenAI 兼容图像生成 API 测试") + print("=" * 60) + print(f"API 地址: {args.base_url}") + print(f"API 密钥: {args.api_key[:8]}...") + print(f"测试提示词: {args.prompt[:50]}...") + + # 创建 OpenAI 客户端 + client = OpenAI( + base_url=f"{args.base_url}/v1", + api_key=args.api_key + ) + + results = [] + + if not args.skip_generation: + # 测试 1: URL 响应格式 + results.append(("URL 响应格式", test_image_generation_url(client, args.prompt))) + + # 测试 2: b64_json 响应格式 + results.append(("b64_json 响应格式", test_image_generation_b64(client, args.prompt))) + + # 测试 4: 响应结构验证 + results.append(("响应结构验证", test_response_structure(client, args.prompt))) + + # 测试 3: 错误处理 + results.append(("错误处理", test_error_handling(client))) + + # 打印总结 + print("\n" + "=" * 60) + print("测试总结") + print("=" * 60) + + passed = 0 + failed = 0 + for name, result in results: + status = "✅ 通过" if result else "❌ 失败" + print(f" {name}: {status}") + if result: + passed += 1 + else: + failed += 1 + + print(f"\n总计: {passed} 通过, {failed} 失败") + + if failed > 0: + sys.exit(1) + else: + print("\n🎉 所有测试通过!") + sys.exit(0) + + +if __name__ == "__main__": + main() diff --git a/src-tauri/src/converter/openai_to_antigravity.rs b/src-tauri/src/converter/openai_to_antigravity.rs index 7404eeea1..ed53ffd9a 100644 --- a/src-tauri/src/converter/openai_to_antigravity.rs +++ b/src-tauri/src/converter/openai_to_antigravity.rs @@ -152,8 +152,9 @@ pub struct AntigravityRequestInner { pub system_instruction: Option, #[serde(skip_serializing_if = "Option::is_none")] pub generation_config: Option, + /// 工具定义 - 支持 Gemini 格式（function_declarations）和 Claude 格式（custom + input_schema） #[serde(skip_serializing_if = "Option::is_none")] - pub tools: Option>, + pub tools: Option, #[serde(skip_serializing_if = "Option::is_none")] pub tool_config: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -251,6 +252,12 @@ fn is_image_generation_model(model: &str) -> bool { model == "gemini-3-pro-image" || model == "gemini-3-pro-image-preview" } +/// 检查是否是 Claude 模型（通过 Antigravity 访问） +/// Claude 模型需要使用不同的工具格式（custom + input_schema） +fn is_claude_model(model: &str) -> bool { + model.starts_with("claude-") || model.contains("claude") +} + // ============================================================================ // 主转换函数 // ============================================================================ @@ -536,6 +543,15 @@ pub fn convert_openai_to_antigravity_with_context( "[ANTIGRAVITY] 图片生成模型 {} 已启用 IMAGE 响应模态", actual_model ); + eprintln!( + "[ANTIGRAVITY] 图片生成模型 {} 已启用 IMAGE 响应模态", + actual_model + ); + } else { + eprintln!( + "[ANTIGRAVITY] 模型 {} 不是图片生成模型，不启用 IMAGE 响应模态", + actual_model + ); } // 处理 reasoning_effort（思维链配置） @@ -573,30 +589,49 @@ pub fn convert_openai_to_antigravity_with_context( } // 转换工具定义 - let tools: Option> = request.tools.as_ref().and_then(|tools| { - let mut function_declarations: Vec = Vec::new(); + // 注意：Antigravity API 统一使用 Gemini 格式（function_declarations） + // Claude 模型在 Antigravity 内部会自动转换 + let tools: Option = request.tools.as_ref().and_then(|tools| { + let is_claude = is_claude_model(actual_model); + + // Gemini 模型使用 function_declarations + parametersJsonSchema + // Claude 模型使用 function_declarations + inputSchema（注意字段名不同） + let mut function_declarations: Vec = Vec::new(); for t in tools { match t { Tool::Function { function } => { - // 转换 parameters -> parametersJsonSchema - let params_schema = function.parameters.as_ref().map(|p| { - let mut schema = clean_parameters(Some(p.clone())).unwrap_or_default(); - // 确保有 type 和 properties - if schema.get("type").is_none() { - schema["type"] = serde_json::json!("object"); - } - if schema.get("properties").is_none() { - schema["properties"] = serde_json::json!({}); - } - schema - }); + let params_schema = function + .parameters + .as_ref() + .map(|p| { + let mut schema = clean_parameters(Some(p.clone())).unwrap_or_default(); + // 确保有 type 和 properties + if schema.get("type").is_none() { + schema["type"] = serde_json::json!("object"); + } + if schema.get("properties").is_none() { + schema["properties"] = serde_json::json!({}); + } + schema + }) + .unwrap_or_else(|| serde_json::json!({"type": "object", "properties": {}})); - function_declarations.push(GeminiFunctionDeclaration { - name: function.name.clone(), - description: function.description.clone(), - parameters_json_schema: params_schema, - }); + if is_claude { + // Claude 模型使用 inputSchema 字段名 + function_declarations.push(serde_json::json!({ + "name": function.name, + "description": function.description.clone().unwrap_or_default(), + "inputSchema": params_schema + })); + } else { + // Gemini 模型使用 parametersJsonSchema 字段名 + function_declarations.push(serde_json::json!({ + "name": function.name, + "description": function.description.clone(), + "parametersJsonSchema": params_schema + })); + } } Tool::WebSearch | Tool::WebSearch20250305 => { // web_search 工具不转换 @@ -607,10 +642,9 @@ pub fn convert_openai_to_antigravity_with_context( if function_declarations.is_empty() { None } else { - Some(vec![GeminiTool { - function_declarations: Some(function_declarations), - google_search: None, - }]) + Some(serde_json::json!([{ + "functionDeclarations": function_declarations + }])) } }); @@ -965,3 +999,617 @@ pub fn convert_antigravity_to_openai_response( response } + +// ============================================================================ +// 图像生成 API 转换函数 +// ============================================================================ + +use crate::models::openai::{ImageData, ImageGenerationRequest, ImageGenerationResponse}; + +/// 图像生成模型名称映射 +/// +/// 注意：Antigravity API 使用内部模型名称 `gemini-3-pro-image`， +/// 而不是用户友好名称 `gemini-3-pro-image-preview`。 +fn image_model_mapping(model: &str) -> &str { + match model { + // OpenAI 兼容模型名 -> Antigravity 内部名称 + "dall-e-3" | "dall-e-2" => "gemini-3-pro-image", + // 用户友好名称 -> 内部名称 + "gemini-3-pro-image-preview" => "gemini-3-pro-image", + _ => model, + } +} + +/// 将 OpenAI 图像生成请求转换为 Antigravity 格式 +/// +/// # 参数 +/// - `request`: OpenAI 图像生成请求 +/// - `project_id`: Antigravity 项目 ID +/// +/// # 返回 +/// Antigravity 格式的请求 JSON +pub fn convert_image_request_to_antigravity( + request: &ImageGenerationRequest, + project_id: &str, +) -> serde_json::Value { + // 模型映射 + let actual_model = image_model_mapping(&request.model); + + // 构建 Gemini 内容结构 + let contents = vec![serde_json::json!({ + "role": "user", + "parts": [{"text": request.prompt}] + })]; + + // 构建生成配置 + let generation_config = serde_json::json!({ + "temperature": 1.0, + "maxOutputTokens": 8096, + "responseModalities": ["TEXT", "IMAGE"], + "candidateCount": request.n + }); + + // 构建安全设置 + let safety_settings = default_safety_settings(); + + // 构建完整请求 + serde_json::json!({ + "project": project_id, + "requestId": format!("img-{}", Uuid::new_v4()), + "request": { + "contents": contents, + "generationConfig": generation_config, + "safetySettings": safety_settings + }, + "model": actual_model, + "userAgent": "antigravity" + }) +} + +/// 将 Antigravity 图像响应转换为 OpenAI 格式 +/// +/// # 参数 +/// - `antigravity_resp`: Antigravity 响应 JSON +/// - `response_format`: 响应格式 ("url" 或 "b64_json") +/// +/// # 返回 +/// OpenAI 格式的图像生成响应，或错误信息 +pub fn convert_antigravity_image_response( + antigravity_resp: &serde_json::Value, + response_format: &str, +) -> Result { + let resp = antigravity_resp.get("response").unwrap_or(antigravity_resp); + + let mut images = Vec::new(); + let mut revised_prompt: Option = None; + + if let Some(candidates) = resp.get("candidates").and_then(|c| c.as_array()) { + for candidate in candidates { + if let Some(parts) = candidate + .get("content") + .and_then(|c| c.get("parts")) + .and_then(|p| p.as_array()) + { + for part in parts { + // 提取文本作为 revised_prompt + if let Some(text) = part.get("text").and_then(|t| t.as_str()) { + if !text.is_empty() { + revised_prompt = Some(text.to_string()); + } + } + + // 提取图像数据 + if let Some(inline_data) = + part.get("inlineData").or_else(|| part.get("inline_data")) + { + if let (Some(data), Some(mime_type)) = ( + inline_data.get("data").and_then(|d| d.as_str()), + inline_data + .get("mimeType") + .or_else(|| inline_data.get("mime_type")) + .and_then(|m| m.as_str()), + ) { + let image_data = if response_format == "b64_json" { + ImageData { + b64_json: Some(data.to_string()), + url: None, + revised_prompt: revised_prompt.clone(), + } + } else { + // 构建 data URL + let data_url = format!("data:{};base64,{}", mime_type, data); + ImageData { + b64_json: None, + url: Some(data_url), + revised_prompt: revised_prompt.clone(), + } + }; + images.push(image_data); + } + } + } + } + } + } + + if images.is_empty() { + return Err("No image generated".to_string()); + } + + Ok(ImageGenerationResponse { + created: chrono::Utc::now().timestamp(), + data: images, + }) +} + +// ============================================================================ +// 图像生成 API 测试 +// ============================================================================ + +#[cfg(test)] +mod image_tests { + use super::*; + + #[test] + fn test_image_model_mapping() { + // OpenAI 兼容模型名映射到内部名称 + assert_eq!(image_model_mapping("dall-e-3"), "gemini-3-pro-image"); + assert_eq!(image_model_mapping("dall-e-2"), "gemini-3-pro-image"); + // 用户友好名称映射到内部名称 + assert_eq!( + image_model_mapping("gemini-3-pro-image-preview"), + "gemini-3-pro-image" + ); + // 内部名称保持不变 + assert_eq!( + image_model_mapping("gemini-3-pro-image"), + "gemini-3-pro-image" + ); + // 其他模型名称透传 + assert_eq!(image_model_mapping("other-model"), "other-model"); + } + + #[test] + fn test_convert_image_request_basic() { + let request = ImageGenerationRequest { + prompt: "A cute cat".to_string(), + model: "gemini-3-pro-image-preview".to_string(), + n: 1, + size: None, + response_format: "url".to_string(), + quality: None, + style: None, + user: None, + }; + + let result = convert_image_request_to_antigravity(&request, "test-project"); + + // 验证基本结构 + assert_eq!(result["project"], "test-project"); + // 模型名应该映射为内部名称 + assert_eq!(result["model"], "gemini-3-pro-image"); + assert!(result["requestId"].as_str().unwrap().starts_with("img-")); + + // 验证内容 + let contents = result["request"]["contents"].as_array().unwrap(); + assert_eq!(contents.len(), 1); + assert_eq!(contents[0]["role"], "user"); + assert_eq!(contents[0]["parts"][0]["text"], "A cute cat"); + + // 验证生成配置 + let gen_config = &result["request"]["generationConfig"]; + let modalities = gen_config["responseModalities"].as_array().unwrap(); + assert!(modalities.contains(&serde_json::json!("TEXT"))); + assert!(modalities.contains(&serde_json::json!("IMAGE"))); + assert_eq!(gen_config["candidateCount"], 1); + } + + #[test] + fn test_convert_image_request_with_n() { + let request = ImageGenerationRequest { + prompt: "A beautiful sunset".to_string(), + model: "dall-e-3".to_string(), + n: 3, + size: Some("1024x1024".to_string()), + response_format: "b64_json".to_string(), + quality: Some("hd".to_string()), + style: Some("vivid".to_string()), + user: None, + }; + + let result = convert_image_request_to_antigravity(&request, "project-123"); + + // dall-e-3 应该映射为内部名称 + assert_eq!(result["model"], "gemini-3-pro-image"); + assert_eq!(result["request"]["generationConfig"]["candidateCount"], 3); + } + + #[test] + fn test_convert_antigravity_image_response_b64_json() { + let antigravity_resp = serde_json::json!({ + "response": { + "candidates": [{ + "content": { + "parts": [ + {"text": "Here is your image"}, + { + "inlineData": { + "mimeType": "image/png", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + } + } + ] + } + }] + } + }); + + let result = convert_antigravity_image_response(&antigravity_resp, "b64_json").unwrap(); + + assert!(result.created > 0); + assert_eq!(result.data.len(), 1); + assert!(result.data[0].b64_json.is_some()); + assert!(result.data[0].url.is_none()); + assert_eq!( + result.data[0].revised_prompt, + Some("Here is your image".to_string()) + ); + } + + #[test] + fn test_convert_antigravity_image_response_url() { + let antigravity_resp = serde_json::json!({ + "response": { + "candidates": [{ + "content": { + "parts": [{ + "inlineData": { + "mimeType": "image/jpeg", + "data": "base64data" + } + }] + } + }] + } + }); + + let result = convert_antigravity_image_response(&antigravity_resp, "url").unwrap(); + + assert!(result.created > 0); + assert_eq!(result.data.len(), 1); + assert!(result.data[0].b64_json.is_none()); + assert!(result.data[0].url.is_some()); + assert_eq!( + result.data[0].url, + Some("data:image/jpeg;base64,base64data".to_string()) + ); + } + + #[test] + fn test_convert_antigravity_image_response_no_image() { + let antigravity_resp = serde_json::json!({ + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "Sorry, I cannot generate that image"}] + } + }] + } + }); + + let result = convert_antigravity_image_response(&antigravity_resp, "url"); + assert!(result.is_err()); + assert_eq!(result.unwrap_err(), "No image generated"); + } + + #[test] + fn test_convert_antigravity_image_response_snake_case() { + // 测试 snake_case 字段名兼容性 + let antigravity_resp = serde_json::json!({ + "response": { + "candidates": [{ + "content": { + "parts": [{ + "inline_data": { + "mime_type": "image/png", + "data": "testdata" + } + }] + } + }] + } + }); + + let result = convert_antigravity_image_response(&antigravity_resp, "b64_json").unwrap(); + assert_eq!(result.data.len(), 1); + assert_eq!(result.data[0].b64_json, Some("testdata".to_string())); + } +} + +// ============================================================================ +// 图像生成 API 属性测试 +// ============================================================================ + +#[cfg(test)] +mod image_property_tests { + use super::*; + use proptest::prelude::*; + + // 生成随机提示词 + fn arb_prompt() -> impl Strategy { + "[a-zA-Z0-9 .,!?]{1,200}".prop_map(|s| s) + } + + // 生成随机模型名称 + fn arb_image_model() -> impl Strategy { + prop_oneof![ + Just("dall-e-3".to_string()), + Just("dall-e-2".to_string()), + Just("gemini-3-pro-image-preview".to_string()), + Just("gemini-3-pro-image".to_string()), + Just("other-model".to_string()), + ] + } + + // 生成随机 n 值 + fn arb_n() -> impl Strategy { + 1u32..5u32 + } + + // 生成随机响应格式 + fn arb_response_format() -> impl Strategy { + prop_oneof![Just("url".to_string()), Just("b64_json".to_string()),] + } + + // 生成随机图像请求 + fn arb_image_request() -> impl Strategy { + ( + arb_prompt(), + arb_image_model(), + arb_n(), + arb_response_format(), + ) + .prop_map( + |(prompt, model, n, response_format)| ImageGenerationRequest { + prompt, + model, + n, + size: None, + response_format, + quality: None, + style: None, + user: None, + }, + ) + } + + // 生成随机 Base64 数据 + fn arb_base64_data() -> impl Strategy { + "[a-zA-Z0-9+/]{10,100}={0,2}".prop_map(|s| s) + } + + // 生成随机 MIME 类型 + fn arb_mime_type() -> impl Strategy { + prop_oneof![ + Just("image/png".to_string()), + Just("image/jpeg".to_string()), + Just("image/gif".to_string()), + Just("image/webp".to_string()), + ] + } + + proptest! { + /// Property 1: Request Conversion Correctness + /// + /// *For any* valid OpenAI image generation request with a non-empty prompt, + /// the converted Antigravity request SHALL: + /// - Contain the prompt text in `request.contents[0].parts[0].text` + /// - Have `responseModalities` set to `["TEXT", "IMAGE"]` + /// - Use the correct mapped model name + /// - Include the specified `n` value in `candidateCount` + /// + /// **Feature: antigravity-image-api, Property 1: Request Conversion Correctness** + /// **Validates: Requirements 1.2, 1.3, 1.4, 2.1, 2.2** + #[test] + fn prop_request_conversion_correctness(request in arb_image_request()) { + let result = convert_image_request_to_antigravity(&request, "test-project"); + + // 验证 prompt 正确传递 + let contents = result["request"]["contents"].as_array().unwrap(); + prop_assert_eq!(contents.len(), 1); + prop_assert_eq!(contents[0]["parts"][0]["text"].as_str().unwrap(), request.prompt.as_str()); + + // 验证 responseModalities 设置正确 + let modalities = result["request"]["generationConfig"]["responseModalities"] + .as_array() + .unwrap(); + prop_assert!(modalities.contains(&serde_json::json!("TEXT"))); + prop_assert!(modalities.contains(&serde_json::json!("IMAGE"))); + + // 验证模型映射正确 + let expected_model = image_model_mapping(&request.model); + prop_assert_eq!(result["model"].as_str().unwrap(), expected_model); + + // 验证 n 值正确传递 + prop_assert_eq!( + result["request"]["generationConfig"]["candidateCount"].as_u64().unwrap(), + request.n as u64 + ); + } + + /// Property 2: Response Format Correctness + /// + /// *For any* Antigravity response containing image data: + /// - WHEN response_format is "b64_json", the output SHALL have `b64_json` field set and `url` field null + /// - WHEN response_format is "url", the output SHALL have `url` field as a valid data URL and `b64_json` field null + /// + /// **Feature: antigravity-image-api, Property 2: Response Format Correctness** + /// **Validates: Requirements 1.6, 1.7, 3.2, 3.3** + #[test] + fn prop_response_format_correctness( + base64_data in arb_base64_data(), + mime_type in arb_mime_type(), + response_format in arb_response_format() + ) { + let antigravity_resp = serde_json::json!({ + "response": { + "candidates": [{ + "content": { + "parts": [{ + "inlineData": { + "mimeType": mime_type.clone(), + "data": base64_data.clone() + } + }] + } + }] + } + }); + + let result = convert_antigravity_image_response(&antigravity_resp, &response_format).unwrap(); + + prop_assert!(result.data.len() >= 1); + + if response_format == "b64_json" { + // b64_json 格式 + prop_assert!(result.data[0].b64_json.is_some()); + prop_assert!(result.data[0].url.is_none()); + prop_assert_eq!(result.data[0].b64_json.as_ref().unwrap(), &base64_data); + } else { + // url 格式 + prop_assert!(result.data[0].b64_json.is_none()); + prop_assert!(result.data[0].url.is_some()); + + // 验证 data URL 格式 + let url = result.data[0].url.as_ref().unwrap(); + let expected_url = format!("data:{};base64,{}", mime_type, base64_data); + prop_assert_eq!(url, &expected_url); + } + } + + /// Property 3: OpenAI Response Structure Compliance + /// + /// *For any* successful image generation response: + /// - The response SHALL have a `created` field with a valid Unix timestamp (positive integer) + /// - The response SHALL have a `data` field that is a non-empty array + /// - Each item in `data` SHALL have either `url` or `b64_json` field (not both) + /// + /// **Feature: antigravity-image-api, Property 3: OpenAI Response Structure Compliance** + /// **Validates: Requirements 3.4, 3.5, 5.3, 5.4** + #[test] + fn prop_openai_response_structure( + base64_data in arb_base64_data(), + mime_type in arb_mime_type(), + response_format in arb_response_format() + ) { + let antigravity_resp = serde_json::json!({ + "response": { + "candidates": [{ + "content": { + "parts": [{ + "inlineData": { + "mimeType": mime_type, + "data": base64_data + } + }] + } + }] + } + }); + + let result = convert_antigravity_image_response(&antigravity_resp, &response_format).unwrap(); + + // 验证 created 是有效时间戳 + prop_assert!(result.created > 0); + + // 验证 data 是非空数组 + prop_assert!(!result.data.is_empty()); + + // 验证每个 item 只有 url 或 b64_json 之一 + for item in &result.data { + let has_url = item.url.is_some(); + let has_b64 = item.b64_json.is_some(); + prop_assert!(has_url != has_b64, "Each item should have exactly one of url or b64_json"); + } + } + + /// Property 4: Model Name Mapping Correctness + /// + /// *For any* model name in the request: + /// - "dall-e-3" SHALL map to "gemini-3-pro-image" + /// - "dall-e-2" SHALL map to "gemini-3-pro-image" + /// - "gemini-3-pro-image-preview" SHALL map to "gemini-3-pro-image" + /// - Other model names SHALL pass through unchanged + /// + /// **Feature: antigravity-image-api, Property 4: Model Name Mapping Correctness** + /// **Validates: Requirements 2.3, 2.4** + #[test] + fn prop_model_name_mapping(model in arb_image_model()) { + let mapped = image_model_mapping(&model); + + match model.as_str() { + "dall-e-3" | "dall-e-2" | "gemini-3-pro-image-preview" => { + prop_assert_eq!(mapped, "gemini-3-pro-image"); + } + other => { + prop_assert_eq!(mapped, other); + } + } + } + + /// Property 5: Image Data Extraction + /// + /// *For any* Antigravity response with `inlineData` containing `data` and `mimeType`: + /// - The converter SHALL successfully extract the Base64 data + /// - The converter SHALL successfully extract the MIME type + /// - If text is present alongside the image, it SHALL be included in `revised_prompt` + /// + /// **Feature: antigravity-image-api, Property 5: Image Data Extraction** + /// **Validates: Requirements 3.1, 3.6** + #[test] + fn prop_image_data_extraction( + base64_data in arb_base64_data(), + mime_type in arb_mime_type(), + text in proptest::option::of("[a-zA-Z0-9 ]{0,100}") + ) { + let mut parts = vec![]; + + // 可选的文本部分 + if let Some(ref t) = text { + if !t.is_empty() { + parts.push(serde_json::json!({"text": t})); + } + } + + // 图像部分 + parts.push(serde_json::json!({ + "inlineData": { + "mimeType": mime_type.clone(), + "data": base64_data.clone() + } + })); + + let antigravity_resp = serde_json::json!({ + "response": { + "candidates": [{ + "content": { + "parts": parts + } + }] + } + }); + + let result = convert_antigravity_image_response(&antigravity_resp, "b64_json").unwrap(); + + // 验证 Base64 数据提取 + prop_assert_eq!(result.data[0].b64_json.as_ref().unwrap(), &base64_data); + + // 验证 revised_prompt + if let Some(ref t) = text { + if !t.is_empty() { + prop_assert_eq!(result.data[0].revised_prompt.as_ref().unwrap(), t); + } + } + } + } +} diff --git a/src-tauri/src/models/openai.rs b/src-tauri/src/models/openai.rs index bef766471..e6454a789 100644 --- a/src-tauri/src/models/openai.rs +++ b/src-tauri/src/models/openai.rs @@ -189,3 +189,82 @@ pub struct ChatCompletionChunk { pub model: String, pub choices: Vec, } + +// ============================================================================ +// 图像生成 API 数据模型 +// ============================================================================ + +/// OpenAI 图像生成请求 +/// +/// 兼容 OpenAI Images API，支持通过 Antigravity 生成图像。 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageGenerationRequest { + /// 图像生成提示词 + pub prompt: String, + + /// 模型名称 (默认: gemini-3-pro-image-preview) + #[serde(default = "default_image_model")] + pub model: String, + + /// 生成图像数量 (默认: 1) + #[serde(default = "default_n")] + pub n: u32, + + /// 图像尺寸 (可选，Antigravity 可能忽略) + #[serde(skip_serializing_if = "Option::is_none")] + pub size: Option, + + /// 响应格式: "url" 或 "b64_json" (默认: "url") + #[serde(default = "default_response_format")] + pub response_format: String, + + /// 图像质量 (可选，Antigravity 可能忽略) + #[serde(skip_serializing_if = "Option::is_none")] + pub quality: Option, + + /// 图像风格 (可选，Antigravity 可能忽略) + #[serde(skip_serializing_if = "Option::is_none")] + pub style: Option, + + /// 用户标识 (可选) + #[serde(skip_serializing_if = "Option::is_none")] + pub user: Option, +} + +fn default_image_model() -> String { + "gemini-3-pro-image-preview".to_string() +} + +fn default_n() -> u32 { + 1 +} + +fn default_response_format() -> String { + "url".to_string() +} + +/// OpenAI 图像生成响应 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageGenerationResponse { + /// 创建时间戳 (Unix epoch seconds) + pub created: i64, + + /// 生成的图像数组 + pub data: Vec, +} + +/// 单个图像数据 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ImageData { + /// Base64 编码的图像数据 (当 response_format="b64_json") + #[serde(skip_serializing_if = "Option::is_none")] + pub b64_json: Option, + + /// 图像 URL (当 response_format="url"，返回 data URL) + #[serde(skip_serializing_if = "Option::is_none")] + pub url: Option, + + /// 修订后的提示词 (如果 Antigravity 返回了文本) + #[serde(skip_serializing_if = "Option::is_none")] + pub revised_prompt: Option, +} diff --git a/src-tauri/src/providers/antigravity.rs b/src-tauri/src/providers/antigravity.rs index 05707403a..c6d5ab3b8 100644 --- a/src-tauri/src/providers/antigravity.rs +++ b/src-tauri/src/providers/antigravity.rs @@ -2007,13 +2007,21 @@ impl StreamingProvider for AntigravityProvider { match result { Ok(resp) => { let status = resp.status(); + eprintln!("[ANTIGRAVITY_STREAM] HTTP 响应状态: {}", status); tracing::info!("[ANTIGRAVITY_STREAM] HTTP 响应状态: {}", status); if status.is_success() { + eprintln!("[ANTIGRAVITY_STREAM] ✓ 流式响应成功建立"); tracing::info!("[ANTIGRAVITY_STREAM] ✓ 流式响应成功建立，返回流"); return Ok(reqwest_stream_to_stream_response(resp)); } else { let body = resp.text().await.unwrap_or_default(); + eprintln!( + "[ANTIGRAVITY_STREAM] ✗ 请求失败\n Base URL: {}\n Status: {}\n Body: {}", + base_url, + status, + &body[..body.len().min(500)] + ); tracing::error!( "[ANTIGRAVITY_STREAM] ✗ 请求失败\n Base URL: {}\n Status: {}\n Body: {}", base_url, @@ -2024,6 +2032,10 @@ impl StreamingProvider for AntigravityProvider { } } Err(e) => { + eprintln!( + "[ANTIGRAVITY_STREAM] ✗ 连接失败\n Base URL: {}\n Error: {}", + base_url, e + ); tracing::error!( "[ANTIGRAVITY_STREAM] ✗ 连接失败\n Base URL: {}\n Error: {}", base_url, diff --git a/src-tauri/src/server/handlers/image_handler.rs b/src-tauri/src/server/handlers/image_handler.rs new file mode 100644 index 000000000..6fd023018 --- /dev/null +++ b/src-tauri/src/server/handlers/image_handler.rs @@ -0,0 +1,352 @@ +//! 图像生成 API 处理器 +//! +//! 实现 OpenAI 兼容的 `/v1/images/generations` 端点， +//! 通过 Antigravity Provider 调用 Gemini 图像生成模型。 +//! +//! # 功能 +//! - 接收 OpenAI 格式的图像生成请求 +//! - 转换为 Antigravity/Gemini 格式 +//! - 调用 Antigravity Provider +//! - 返回 OpenAI 格式的响应 +//! +//! # 需求覆盖 +//! - 需求 1.1: 实现 `/v1/images/generations` 端点 +//! - 需求 4.1: 验证请求参数 +//! - 需求 4.2: 获取 Antigravity 凭证 +//! - 需求 4.3: 调用 Antigravity Provider +//! - 需求 4.4: 转换响应格式 + +use axum::{ + extract::State, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; + +use crate::converter::openai_to_antigravity::{ + convert_antigravity_image_response, convert_image_request_to_antigravity, +}; +use crate::models::openai::ImageGenerationRequest; +use crate::models::provider_pool_model::CredentialData; +use crate::providers::AntigravityProvider; +use crate::server::handlers::verify_api_key; +use crate::server::AppState; + +/// 处理图像生成请求 +/// +/// # 端点 +/// `POST /v1/images/generations` +/// +/// # 请求格式 +/// ```json +/// { +/// "prompt": "A cute cat", +/// "model": "dall-e-3", +/// "n": 1, +/// "size": "1024x1024", +/// "response_format": "url" +/// } +/// ``` +/// +/// # 响应格式 +/// ```json +/// { +/// "created": 1234567890, +/// "data": [ +/// { +/// "url": "data:image/png;base64,...", +/// "revised_prompt": "A cute fluffy cat" +/// } +/// ] +/// } +/// ``` +pub async fn handle_image_generation( + State(state): State, + headers: HeaderMap, + Json(request): Json, +) -> Response { + // 验证 API Key + if let Err(e) = verify_api_key(&headers, &state.api_key).await { + return e.into_response(); + } + + // 验证请求参数 + if request.prompt.trim().is_empty() { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": { + "message": "prompt is required and cannot be empty", + "type": "invalid_request_error", + "code": "invalid_prompt" + } + })), + ) + .into_response(); + } + + // 记录请求日志 + state.logs.write().await.add( + "info", + &format!( + "[IMAGE] 收到图像生成请求: model={}, prompt={}, n={}, response_format={}", + request.model, + if request.prompt.len() > 50 { + format!("{}...", &request.prompt[..50]) + } else { + request.prompt.clone() + }, + request.n, + request.response_format + ), + ); + + // 获取 Antigravity 凭证 + let db = match &state.db { + Some(db) => db, + None => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": { + "message": "Database not available", + "type": "server_error" + } + })), + ) + .into_response(); + } + }; + + // 从凭证池获取 Antigravity 凭证 + let credential = match state + .pool_service + .select_credential(db, "antigravity", None) + { + Ok(Some(cred)) => cred, + Ok(None) => { + state + .logs + .write() + .await + .add("error", "[IMAGE] 没有可用的 Antigravity 凭证"); + return ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "error": { + "message": "No Antigravity credentials available for image generation", + "type": "server_error", + "code": "no_credentials" + } + })), + ) + .into_response(); + } + Err(e) => { + state + .logs + .write() + .await + .add("error", &format!("[IMAGE] 获取凭证失败: {}", e)); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": { + "message": format!("Failed to get credentials: {}", e), + "type": "server_error" + } + })), + ) + .into_response(); + } + }; + + // 提取 Antigravity 凭证信息 + let (creds_file_path, project_id) = match &credential.credential { + CredentialData::AntigravityOAuth { + creds_file_path, + project_id, + } => (creds_file_path.clone(), project_id.clone()), + _ => { + state + .logs + .write() + .await + .add("error", "[IMAGE] 选中的凭证不是 Antigravity 类型"); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": { + "message": "Selected credential is not Antigravity type", + "type": "server_error" + } + })), + ) + .into_response(); + } + }; + + // 创建 Antigravity Provider + let mut antigravity = AntigravityProvider::new(); + if let Err(e) = antigravity + .load_credentials_from_path(&creds_file_path) + .await + { + let _ = state.pool_service.mark_unhealthy( + db, + &credential.uuid, + Some(&format!("Failed to load credentials: {}", e)), + ); + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": { + "message": format!("Failed to load Antigravity credentials: {}", e), + "type": "server_error" + } + })), + ) + .into_response(); + } + + // 验证并刷新 Token + let validation_result = antigravity.validate_token(); + if validation_result.needs_refresh() { + tracing::info!("[IMAGE] Token 需要刷新，开始刷新..."); + if let Err(refresh_error) = antigravity.refresh_token_with_retry(3).await { + tracing::error!("[IMAGE] Token 刷新失败: {:?}", refresh_error); + let _ = state.pool_service.mark_unhealthy_with_details( + db, + &credential.uuid, + &refresh_error, + ); + let (status, message) = if refresh_error.requires_reauth() { + (StatusCode::UNAUTHORIZED, refresh_error.user_message()) + } else { + ( + StatusCode::INTERNAL_SERVER_ERROR, + refresh_error.user_message(), + ) + }; + return ( + status, + Json(serde_json::json!({ + "error": { + "message": message, + "type": "authentication_error" + } + })), + ) + .into_response(); + } + } + + // 设置项目 ID + if let Some(pid) = project_id { + antigravity.project_id = Some(pid); + } else if let Err(e) = antigravity.discover_project().await { + tracing::warn!("[IMAGE] Failed to discover project: {}", e); + } + + let proj_id = antigravity.project_id.clone().unwrap_or_default(); + + // 转换请求为 Antigravity 格式 + let antigravity_request = convert_image_request_to_antigravity(&request, &proj_id); + + state.logs.write().await.add( + "debug", + &format!( + "[IMAGE] Antigravity 请求: model={}", + antigravity_request["model"].as_str().unwrap_or("unknown") + ), + ); + + // 调用 Antigravity API - 直接使用 call_api 而不是 generate_content + // 因为 generate_content 内部的 to_gemini_response 会丢失嵌套在 response 字段下的数据 + let model = antigravity_request["model"] + .as_str() + .unwrap_or("gemini-3-pro-image-preview"); + + eprintln!("[IMAGE] 调用 Antigravity API: model={}", model); + eprintln!( + "[IMAGE] 请求内容: {}", + serde_json::to_string_pretty(&antigravity_request).unwrap_or_default() + ); + + match antigravity + .call_api("generateContent", &antigravity_request) + .await + { + Ok(resp) => { + // 调试：打印原始响应 + eprintln!( + "[IMAGE] Antigravity 原始响应: {}", + serde_json::to_string_pretty(&resp).unwrap_or_default() + ); + state.logs.write().await.add( + "debug", + &format!( + "[IMAGE] Antigravity 原始响应: {}", + serde_json::to_string(&resp).unwrap_or_default() + ), + ); + + // 转换响应为 OpenAI 格式 + match convert_antigravity_image_response(&resp, &request.response_format) { + Ok(image_response) => { + // 记录成功 + let _ = state + .pool_service + .mark_healthy(db, &credential.uuid, Some(model)); + let _ = state.pool_service.record_usage(db, &credential.uuid); + + state.logs.write().await.add( + "info", + &format!("[IMAGE] 图像生成成功: {} 张图片", image_response.data.len()), + ); + + (StatusCode::OK, Json(image_response)).into_response() + } + Err(e) => { + state + .logs + .write() + .await + .add("error", &format!("[IMAGE] 响应转换失败: {}", e)); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": { + "message": e, + "type": "server_error", + "code": "image_generation_failed" + } + })), + ) + .into_response() + } + } + } + Err(e) => { + let _ = state + .pool_service + .mark_unhealthy(db, &credential.uuid, Some(&e.to_string())); + state + .logs + .write() + .await + .add("error", &format!("[IMAGE] Antigravity API 调用失败: {}", e)); + ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "error": { + "message": format!("Image generation failed: {}", e), + "type": "server_error", + "code": "api_error" + } + })), + ) + .into_response() + } + } +} diff --git a/src-tauri/src/server/handlers/mod.rs b/src-tauri/src/server/handlers/mod.rs index da731f254..996e0f029 100644 --- a/src-tauri/src/server/handlers/mod.rs +++ b/src-tauri/src/server/handlers/mod.rs @@ -4,6 +4,7 @@ pub mod api; pub mod credentials_api; +pub mod image_handler; pub mod kiro_credential; pub mod management; pub mod provider_calls; @@ -11,6 +12,7 @@ pub mod websocket; pub use api::*; pub use credentials_api::*; +pub use image_handler::*; pub use kiro_credential::*; pub use management::*; pub use provider_calls::*; diff --git a/src-tauri/src/server/mod.rs b/src-tauri/src/server/mod.rs index 05e71c52c..c3f4e5cbc 100644 --- a/src-tauri/src/server/mod.rs +++ b/src-tauri/src/server/mod.rs @@ -871,6 +871,11 @@ async fn run_server( .route("/v1/chat/completions", post(handlers::chat_completions)) .route("/v1/messages", post(handlers::anthropic_messages)) .route("/v1/messages/count_tokens", post(count_tokens)) + // 图像生成 API 路由 + .route( + "/v1/images/generations", + post(handlers::handle_image_generation), + ) // Gemini 原生协议路由 .route("/v1/gemini/*path", post(gemini_generate_content)) // WebSocket 路由