opensource/proxycast
3
0
Fork 0
mirror of https://github.com/aiclientproxy/proxycast.git synced 2026-09-24 23:10:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 添加 Antigravity 图像生成 API (OpenAI 兼容)

Browse Source 
- 实现 /v1/images/generations 端点
- 添加 ImageGenerationRequest/Response 数据模型
- 实现 OpenAI 到 Antigravity 请求/响应转换器
- 支持 dall-e-3/dall-e-2 模型名映射到 gemini-3-pro-image
- 支持 b64_json 和 url 两种响应格式
- 添加 Python SDK 集成测试脚本
This commit is contained in:
liubu
2026-01-03 23:34:03 +08:00
parent a793de37d9
commit 3ff24e060e
7 changed files with 1491 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/test_image_api.py
Executable
+369
View File
@@ -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()
src-tauri/src/converter/openai_to_antigravity.rs
+672 -24
View File
@@ -152,8 +152,9 @@ pub struct AntigravityRequestInner {
pub system_instruction: Option<GeminiContent>,
#[serde(skip_serializing_if = "Option::is_none")]
pub generation_config: Option<GeminiGenerationConfig>,
/// 工具定义 - 支持 Gemini 格式（function_declarations）和 Claude 格式（custom + input_schema）
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<GeminiTool>>,
pub tools: Option<serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_config: Option<serde_json::Value>,
#[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<Vec<GeminiTool>> = request.tools.as_ref().and_then(|tools| {
let mut function_declarations: Vec<GeminiFunctionDeclaration> = Vec::new();
// 注意：Antigravity API 统一使用 Gemini 格式（function_declarations）
// Claude 模型在 Antigravity 内部会自动转换
let tools: Option<serde_json::Value> = 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<serde_json::Value> = 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<ImageGenerationResponse, String> {
let resp = antigravity_resp.get("response").unwrap_or(antigravity_resp);
let mut images = Vec::new();
let mut revised_prompt: Option<String> = 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<Value = String> {
"[a-zA-Z0-9 .,!?]{1,200}".prop_map(|s| s)
}
// 生成随机模型名称
fn arb_image_model() -> impl Strategy<Value = String> {
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<Value = u32> {
1u32..5u32
}
// 生成随机响应格式
fn arb_response_format() -> impl Strategy<Value = String> {
prop_oneof![Just("url".to_string()), Just("b64_json".to_string()),]
}
// 生成随机图像请求
fn arb_image_request() -> impl Strategy<Value = ImageGenerationRequest> {
(
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<Value = String> {
"[a-zA-Z0-9+/]{10,100}={0,2}".prop_map(|s| s)
}
// 生成随机 MIME 类型
fn arb_mime_type() -> impl Strategy<Value = String> {
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);
}
}
}
}
}
src-tauri/src/models/openai.rs
+79
View File
@@ -189,3 +189,82 @@ pub struct ChatCompletionChunk {
pub model: String,
pub choices: Vec<StreamChoice>,
}
// ============================================================================
// 图像生成 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<String>,
/// 响应格式: "url" 或 "b64_json" (默认: "url")
#[serde(default = "default_response_format")]
pub response_format: String,
/// 图像质量 (可选，Antigravity 可能忽略)
#[serde(skip_serializing_if = "Option::is_none")]
pub quality: Option<String>,
/// 图像风格 (可选，Antigravity 可能忽略)
#[serde(skip_serializing_if = "Option::is_none")]
pub style: Option<String>,
/// 用户标识 (可选)
#[serde(skip_serializing_if = "Option::is_none")]
pub user: Option<String>,
}
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<ImageData>,
}
/// 单个图像数据
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageData {
/// Base64 编码的图像数据 (当 response_format="b64_json")
#[serde(skip_serializing_if = "Option::is_none")]
pub b64_json: Option<String>,
/// 图像 URL (当 response_format="url"，返回 data URL)
#[serde(skip_serializing_if = "Option::is_none")]
pub url: Option<String>,
/// 修订后的提示词 (如果 Antigravity 返回了文本)
#[serde(skip_serializing_if = "Option::is_none")]
pub revised_prompt: Option<String>,
}
src-tauri/src/providers/antigravity.rs
+12
View File
@@ -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,
src-tauri/src/server/handlers/image_handler.rs
+352
View File
@@ -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<AppState>,
headers: HeaderMap,
Json(request): Json<ImageGenerationRequest>,
) -> 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()
}
}
}
src-tauri/src/server/handlers/mod.rs
+2
View File
@@ -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::*;
src-tauri/src/server/mod.rs
+5
View File
@@ -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 路由