tamar-model-client 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

tamar_model_client/core/request_builder.py

@@ -17,6 +17,7 @@ from ..schemas.inputs import (
     UserContext,
     GoogleGenAiInput,
     GoogleVertexAIImagesInput,
+    GoogleGenAIImagesInput,
     OpenAIResponsesInput,
     OpenAIChatCompletionsInput,
     OpenAIImagesInput,
@@ -32,7 +33,7 @@ class RequestBuilder:
     负责将高级的 ModelRequest 对象转换为 gRPC 协议所需的请求对象，
     包括参数验证、序列化和提供商特定的字段处理。
     """
-    
+
     @staticmethod
     def get_allowed_fields(provider: ProviderType, invoke_type: InvokeType) -> Set[str]:
         """
@@ -61,11 +62,13 @@ class RequestBuilder:
                 return set(OpenAIImagesInput.model_fields.keys())
             case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.IMAGE_EDIT_GENERATION):
                 return set(OpenAIImagesEditInput.model_fields.keys())
+            case (ProviderType.GOOGLE, InvokeType.IMAGE_GENERATION_GENAI):
+                return set(GoogleGenAIImagesInput.model_fields.keys())
             case _:
                 raise ValueError(
                     f"Unsupported provider/invoke_type combination: {provider} + {invoke_type}"
                 )
-    
+
     @staticmethod
     def build_grpc_extra_fields(model_request: ModelRequest) -> Dict[str, Any]:
         """
@@ -88,31 +91,31 @@ class RequestBuilder:
                 model_request.provider,
                 model_request.invoke_type
             )
-            
+
             # 将 ModelRequest 转换为字典，只包含已设置的字段
             model_request_dict = model_request.model_dump(exclude_unset=True)
-            
+
             # 构建 gRPC 请求参数
             grpc_request_kwargs = {}
             for field in allowed_fields:
                 if field in model_request_dict:
                     value = model_request_dict[field]
-                    
+
                     # 跳过无效的值
                     if not is_effective_value(value):
                         continue
-                    
+
                     # 序列化不支持的类型
                     grpc_request_kwargs[field] = serialize_value(value)
-            
+
             # 清理序列化后的参数中的 None 值
             grpc_request_kwargs = remove_none_from_dict(grpc_request_kwargs)
-            
+
             return grpc_request_kwargs
-            
+
         except Exception as e:
             raise ValueError(f"构建请求失败: {str(e)}") from e
-    
+
     @staticmethod
     def build_single_request(model_request: ModelRequest) -> model_service_pb2.ModelRequestItem:
         """
@@ -129,7 +132,7 @@ class RequestBuilder:
         """
         # 构建额外字段
         extra_fields = RequestBuilder.build_grpc_extra_fields(model_request)
-        
+
         # 创建 gRPC 请求对象
         return model_service_pb2.ModelRequestItem(
             provider=model_request.provider.value,
@@ -141,11 +144,11 @@ class RequestBuilder:
             client_type=model_request.user_context.client_type or "",
             extra=extra_fields
         )
-    
+
     @staticmethod
     def build_batch_request_item(
-        batch_item: "BatchModelRequestItem", 
-        user_context: "UserContext"
+            batch_item: "BatchModelRequestItem",
+            user_context: "UserContext"
     ) -> model_service_pb2.ModelRequestItem:
         """
         构建批量请求中的单个项目
@@ -159,7 +162,7 @@ class RequestBuilder:
         """
         # 构建额外字段
         extra_fields = RequestBuilder.build_grpc_extra_fields(batch_item)
-        
+
         # 添加 custom_id 如果存在
         if hasattr(batch_item, 'custom_id') and batch_item.custom_id:
             request_item = model_service_pb2.ModelRequestItem(
@@ -184,13 +187,13 @@ class RequestBuilder:
                 client_type=user_context.client_type or "",
                 extra=extra_fields
             )
-            
+
         # 添加 priority 如果存在
         if hasattr(batch_item, 'priority') and batch_item.priority is not None:
             request_item.priority = batch_item.priority
-            
+
         return request_item
-    
+
     @staticmethod
     def build_batch_request(batch_request: BatchModelRequest) -> model_service_pb2.ModelRequest:
         """
@@ -206,16 +209,16 @@ class RequestBuilder:
             ValueError: 当构建请求失败时
         """
         items = []
-        
+
         for batch_item in batch_request.items:
             # 为每个请求项构建 gRPC 对象，传入 user_context
             request_item = RequestBuilder.build_batch_request_item(
-                batch_item, 
+                batch_item,
                 batch_request.user_context
             )
             items.append(request_item)
-        
+
         # 创建批量请求对象
         return model_service_pb2.ModelRequest(
             items=items
-        )
+        )

tamar_model_client/enums/invoke.py

@@ -9,3 +9,4 @@ class InvokeType(str, Enum):
     GENERATION = "generation"  # 生成类，默认的值
     IMAGE_GENERATION = "image-generation"
     IMAGE_EDIT_GENERATION = "image-edit-generation"
+    IMAGE_GENERATION_GENAI = "image-generation-genai"  # GenAI SDK图像生成

tamar_model_client/schemas/inputs.py

@@ -59,6 +59,16 @@ class GoogleVertexAIImagesInput(BaseModel):
     }
 
 
+class GoogleGenAIImagesInput(BaseModel):
+    model: str
+    prompt: str
+    config: Optional[types.GenerateImagesConfigOrDict] = None
+
+    model_config = {
+        "arbitrary_types_allowed": True
+    }
+
+
 class OpenAIResponsesInput(BaseModel):
     background: Optional[bool] | NotGiven = NOT_GIVEN
     include: Optional[List[ResponseIncludable]] | NotGiven = NOT_GIVEN
@@ -264,7 +274,7 @@ class ModelRequestInput(BaseRequest):
 
     # Google GenAI Input
     contents: Optional[Union[types.ContentListUnion, types.ContentListUnionDict]] = None
-    config: Optional[types.GenerateContentConfigOrDict] = None
+    config: Optional[Union[types.GenerateContentConfigOrDict, types.GenerateImagesConfigOrDict]] = None
 
     # Images（OpenAI Images / Images Edit / Google Vertex Images 合并）
     image: Optional[Union[FileTypes, List[FileTypes]]] = None

tamar_model_client/utils.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from typing import Any
 import os, mimetypes
 
+
 def convert_file_field(value: Any) -> Any:
     def is_file_like(obj):
         return hasattr(obj, "read") and callable(obj.read)
@@ -55,7 +56,7 @@ def validate_fields_by_provider_and_invoke_type(
     """
     from tamar_model_client.enums import ProviderType, InvokeType
     from tamar_model_client.schemas.inputs import GoogleGenAiInput, OpenAIResponsesInput, OpenAIChatCompletionsInput, \
-        OpenAIImagesInput, OpenAIImagesEditInput, GoogleVertexAIImagesInput
+        OpenAIImagesInput, OpenAIImagesEditInput, GoogleVertexAIImagesInput, GoogleGenAIImagesInput
 
     google_allowed = extra_allowed_fields | set(GoogleGenAiInput.model_fields)
     openai_responses_allowed = extra_allowed_fields | set(OpenAIResponsesInput.model_fields)
@@ -63,9 +64,11 @@ def validate_fields_by_provider_and_invoke_type(
     openai_images_allowed = extra_allowed_fields | set(OpenAIImagesInput.model_fields)
     openai_images_edit_allowed = extra_allowed_fields | set(OpenAIImagesEditInput.model_fields)
     google_vertexai_images_allowed = extra_allowed_fields | set(GoogleVertexAIImagesInput.model_fields)
+    google_genai_images_allowed = extra_allowed_fields | set(GoogleGenAIImagesInput.model_fields)
 
     google_required = {"model", "contents"}
     google_vertex_required = {"model", "prompt"}
+    google_genai_images_required = {"model", "prompt"}
     openai_resp_required = {"input", "model"}
     openai_chat_required = {"messages", "model"}
     openai_img_required = {"prompt"}
@@ -90,6 +93,9 @@ def validate_fields_by_provider_and_invoke_type(
         case ((ProviderType.OPENAI | ProviderType.AZURE), InvokeType.IMAGE_EDIT_GENERATION):
             allowed = openai_images_edit_allowed
             required = openai_edit_required
+        case (ProviderType.GOOGLE, InvokeType.IMAGE_GENERATION_GENAI):
+            allowed = google_genai_images_allowed
+            required = google_genai_images_required
         case _:
             raise ValueError(f"Unsupported provider/invoke_type: {instance.provider} + {instance.invoke_type}")
 

{tamar_model_client-0.2.4.dist-info → tamar_model_client-0.2.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tamar-model-client
-Version: 0.2.4
+Version: 0.2.5
 Summary: A Python SDK for interacting with the Model Manager gRPC service
 Home-page: http://gitlab.tamaredge.top/project-tap/AgentOS/model-manager-client
 Author: Oscar Ou
@@ -54,7 +54,7 @@ Dynamic: summary
 
 ### 🔌 多服务商支持
 - **OpenAI** (GPT-3.5/4, DALL-E)
-- **Google** (Gemini - AI Studio & Vertex AI)
+- **Google** (Gemini - AI Studio & Vertex AI, Imagen 图像生成)
 - **Azure OpenAI** (企业级部署)
 - **Anthropic** (Claude)
 - **DeepSeek** (深度求索)
@@ -114,7 +114,7 @@ tamar_model_client/
 │   └── outputs.py            # 响应模型（ModelResponse, Usage）
 ├── 📁 enums/                  # 枚举定义
 │   ├── providers.py          # AI 服务商（OpenAI, Google, Azure...）
-│   ├── invoke.py             # 调用类型（generation, images...）
+│   ├── invoke.py             # 调用类型（generation, images, image-generation-genai...）
 │   └── channel.py            # 服务通道（openai, vertexai...）
 ├── 📁 core/                   # 核心功能模块
 │   ├── base_client.py        # 客户端基类（熔断、降级、配置）
@@ -288,6 +288,37 @@ else:
     print(f"响应: {vertex_response.content}")
     if vertex_response.usage:
         print(f"Token 使用情况: {vertex_response.usage}")
+
+# Google GenAI 图像生成示例
+from google.genai import types
+
+genai_image_request = ModelRequest(
+    provider=ProviderType.GOOGLE,  # 选择 Google 作为提供商
+    channel=Channel.AI_STUDIO,  # 使用 AI Studio 渠道
+    invoke_type=InvokeType.IMAGE_GENERATION_GENAI,  # 使用 GenAI 图像生成调用类型
+    model="imagen-3.0-generate-001",  # 指定图像生成模型
+    prompt="一只可爱的小猫在花园里玩耍",  # 图像描述提示词
+    user_context=UserContext(
+        user_id="test_user",
+        org_id="test_org",
+        client_type="python-sdk"
+    ),
+    # 使用 Google GenAI 类型构建配置
+    config=types.GenerateImagesConfig(
+        number_of_images=1,
+        aspect_ratio="1:1",
+        safety_filter_level="block_some"
+    )
+)
+
+# 发送图像生成请求并获取响应
+image_response = client.invoke(genai_image_request)
+if image_response.error:
+    print(f"错误: {image_response.error}")
+else:
+    print(f"图像生成成功: {image_response.content}")
+    if image_response.usage:
+        print(f"使用情况: {image_response.usage}")
 ```
 
 ### Azure OpenAI 调用示例
@@ -488,6 +519,78 @@ async def batch_example():
 asyncio.run(batch_example())
 ```
 
+### 图像生成调用示例
+
+支持 OpenAI DALL-E、Google Vertex AI 和 Google GenAI 图像生成：
+
+```python
+from tamar_model_client import TamarModelClient
+from tamar_model_client.schemas import ModelRequest, UserContext
+from tamar_model_client.enums import ProviderType, InvokeType, Channel
+
+client = TamarModelClient()
+
+# OpenAI DALL-E 图像生成
+openai_image_request = ModelRequest(
+    provider=ProviderType.OPENAI,
+    channel=Channel.OPENAI,
+    invoke_type=InvokeType.IMAGE_GENERATION,
+    model="dall-e-3",
+    prompt="一只穿着西装的猫在办公室里工作",
+    user_context=UserContext(
+        user_id="test_user",
+        org_id="test_org",
+        client_type="python-sdk"
+    ),
+    size="1024x1024",
+    quality="hd",
+    n=1
+)
+
+# Google Vertex AI 图像生成
+vertex_image_request = ModelRequest(
+    provider=ProviderType.GOOGLE,
+    channel=Channel.VERTEXAI,
+    invoke_type=InvokeType.IMAGE_GENERATION,
+    model="imagegeneration@006",
+    prompt="一座美丽的山峰在日出时分",
+    user_context=UserContext(
+        user_id="test_user",
+        org_id="test_org",
+        client_type="python-sdk"
+    ),
+    number_of_images=1,
+    aspect_ratio="1:1",
+    safety_filter_level="block_some"
+)
+
+# Google GenAI 图像生成（新增功能）
+genai_image_request = ModelRequest(
+    provider=ProviderType.GOOGLE,
+    channel=Channel.AI_STUDIO,
+    invoke_type=InvokeType.IMAGE_GENERATION_GENAI,  # 新增的调用类型
+    model="imagen-3.0-generate-001",
+    prompt="科幻风格的城市夜景，霓虹灯闪烁",
+    user_context=UserContext(
+        user_id="test_user",
+        org_id="test_org",
+        client_type="python-sdk"
+    ),
+    config=types.GenerateImagesConfig(
+        number_of_images=1,
+        aspect_ratio="16:9"
+    )
+)
+
+# 发送请求
+for request in [openai_image_request, vertex_image_request, genai_image_request]:
+    response = client.invoke(request)
+    if response.error:
+        print(f"图像生成失败: {response.error}")
+    else:
+        print(f"图像生成成功: {response.content}")
+```
+
 ### 文件输入示例
 
 支持处理图像等文件输入（需使用支持多模态的模型，如 gemini-2.0-flash）：

{tamar_model_client-0.2.4.dist-info → tamar_model_client-0.2.5.dist-info}/RECORD

@@ -7,32 +7,32 @@ tamar_model_client/exceptions.py,sha256=EOr4JMYI7hVszRvNYJ1JqsUNpVmd16T2KpJ0MkFT
 tamar_model_client/json_formatter.py,sha256=XT8XPMKKM2M22tuYR2e1rvWHcpz3UD9iLLgGPsGOjCI,2410
 tamar_model_client/logging_icons.py,sha256=MRTZ1Xvkep9ce_jdltj54_XZUXvIpQ95soRNmLdJ4qw,1837
 tamar_model_client/sync_client.py,sha256=FbyjuyDRiXklSS_l5h5fwNxvABI-hpLGiIWAXqhPHoI,53760
-tamar_model_client/utils.py,sha256=Kn6pFz9GEC96H4eejEax66AkzvsrXI3WCSDtgDjnVTI,5238
+tamar_model_client/utils.py,sha256=9gJm71UuQhyyBCgo6gvMjv74xepOlw6AiwuSzea2CL0,5595
 tamar_model_client/core/__init__.py,sha256=RMiZjV1S4csWPLxB_JfdOea8fYPz97Oj3humQSBw1OI,1054
 tamar_model_client/core/base_client.py,sha256=spb4zjDuPczqnXNlDcIq_bDQ09TOpxeeuX7IxpTS_38,13859
 tamar_model_client/core/http_fallback.py,sha256=2N7-N_TZrtffDjuv9s3-CD8Xy7qw9AuI5xeWGUnGQ0w,22217
 tamar_model_client/core/logging_setup.py,sha256=-MXzTR4Ax50H16cbq1jCXbxgayf5fZ0U3o0--fMmxD8,6692
-tamar_model_client/core/request_builder.py,sha256=yi8iy2Ps2m4d1YwIFiQLRxTvxQxgEGV576aXnNYRl7E,8507
+tamar_model_client/core/request_builder.py,sha256=aplTEXGgeipn-dRCdUptHYWkT9c4zjKmbmI8Ckbv_sM,8516
 tamar_model_client/core/request_id_manager.py,sha256=S-Mliaby9zN_bx-B85FvVnttal-w0skkjy2ZvWoQ5vw,3689
 tamar_model_client/core/response_handler.py,sha256=_q5galAT0_RaUT5C_yZsjg-9VnT9CBjmIASOt28BUmQ,4616
 tamar_model_client/core/utils.py,sha256=AcbsGfNQEaZLYI4OZJs-BdmJgxAoLUC5LFoiYmji820,5875
 tamar_model_client/enums/__init__.py,sha256=3cYYn8ztNGBa_pI_5JGRVYf2QX8fkBVWdjID1PLvoBQ,182
 tamar_model_client/enums/channel.py,sha256=wCzX579nNpTtwzGeS6S3Ls0UzVAgsOlfy4fXMzQTCAw,199
-tamar_model_client/enums/invoke.py,sha256=Up87myAg4-0SDJV5a82ggPDpYHSLEtIco8BF_5Ph1nY,322
+tamar_model_client/enums/invoke.py,sha256=4GLwmUnX9Arwvu-lfFBtYr7aLTB7lZgTQvZEJKUGqY4,402
 tamar_model_client/enums/providers.py,sha256=L_bX75K6KnWURoFizoitZ1Ybza7bmYDqXecNzNpgIrI,165
 tamar_model_client/generated/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tamar_model_client/generated/model_service_pb2.py,sha256=RI6wNSmgmylzWPedFfPxx938UzS7kcPR58YTzYshcL8,3066
 tamar_model_client/generated/model_service_pb2_grpc.py,sha256=k4tIbp3XBxdyuOVR18Ung_4SUryONB51UYf_uUEl6V4,5145
 tamar_model_client/schemas/__init__.py,sha256=AxuI-TcvA4OMTj2FtK4wAItvz9LrK_293pu3cmMLE7k,394
-tamar_model_client/schemas/inputs.py,sha256=ilSZxcnXubX-yndz6X7y_mZnx19f5IGmWoKPSm0L_R8,17069
+tamar_model_client/schemas/inputs.py,sha256=5ZjuWhlIY_geFGL6cnqk2v1ZQ7algyRD2jJ6NiShVt4,17324
 tamar_model_client/schemas/outputs.py,sha256=M_fcqUtXPJnfiLabHlyA8BorlC5pYkf5KLjXO1ysKIQ,1031
 tests/__init__.py,sha256=kbmImddLDwdqlkkmkyKtl4bQy_ipe-R8eskpaBylU9w,38
 tests/stream_hanging_analysis.py,sha256=W3W48IhQbNAR6-xvMpoWZvnWOnr56CTaH4-aORNBuD4,14807
 tests/test_circuit_breaker.py,sha256=nhEBnyXFjIYjRWlUdu7Z9PnPq48ypbBK6fxN6deHedw,12172
-tests/test_google_azure_final.py,sha256=Cx2lfnoj48_7pUjpCYbrx6OLJF4cI79McV24_EYt_8s,55093
+tests/test_google_azure_final.py,sha256=YnTl5_e-C_lay2e0_0z35rpIxRkLs1tTsvcQ3Qww_RU,67077
 tests/test_logging_issue.py,sha256=JTMbotfHpAEPMBj73pOwxPn-Zn4QVQJX6scMz48FRDQ,2427
 tests/test_simple.py,sha256=Xf0U-J9_xn_LzUsmYu06suK0_7DrPeko8OHoHldsNxE,7169
-tamar_model_client-0.2.4.dist-info/METADATA,sha256=YN-OEy64kB_c1f9sTYRxciIdbn5KSAIeRXECpR7xKLU,41309
-tamar_model_client-0.2.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tamar_model_client-0.2.4.dist-info/top_level.txt,sha256=f1I-S8iWN-cgv4gB8gxRg9jJOTJMumvm4oGKVPfGg6A,25
-tamar_model_client-0.2.4.dist-info/RECORD,,
+tamar_model_client-0.2.5.dist-info/METADATA,sha256=d_AzgzHtqfzXlVVHLZU88m3LigTm1oE0Jc_8NIYPggQ,44621
+tamar_model_client-0.2.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tamar_model_client-0.2.5.dist-info/top_level.txt,sha256=f1I-S8iWN-cgv4gB8gxRg9jJOTJMumvm4oGKVPfGg6A,25
+tamar_model_client-0.2.5.dist-info/RECORD,,

tests/test_google_azure_final.py

@@ -35,6 +35,7 @@ try:
     from tamar_model_client import TamarModelClient, AsyncTamarModelClient
     from tamar_model_client.schemas import ModelRequest, UserContext
     from tamar_model_client.enums import ProviderType, InvokeType, Channel
+    from google.genai import types
 
     # 为了调试，临时启用 SDK 的日志输出
     # 注意：这会输出 JSON 格式的日志
@@ -145,6 +146,70 @@ def test_azure_openai():
         print(f"❌ Azure OpenAI 失败: {str(e)}")
 
 
+def test_google_genai_image_generation():
+    """测试 Google GenAI 图像生成"""
+    print("\n🎨 测试 Google GenAI 图像生成...")
+
+    try:
+        client = TamarModelClient()
+
+        request = ModelRequest(
+            provider=ProviderType.GOOGLE,
+            channel=Channel.AI_STUDIO,
+            invoke_type=InvokeType.IMAGE_GENERATION_GENAI,
+            model="imagen-3.0-generate-002",
+            prompt="一只可爱的小猫咪在花园里玩耍，阳光透过树叶洒下斑驳的光影",
+            user_context=UserContext(
+                user_id="test_user",
+                org_id="test_org",
+                client_type="test_client"
+            )
+        )
+
+        response = client.invoke(request)
+        print(f"✅ Google GenAI 图像生成成功")
+        print(f"   响应类型: {type(response)}")
+        if response.content:
+            print(f"   图像内容长度: {len(str(response.content))}")
+        else:
+            print(f"   响应内容: {str(response)[:200]}...")
+
+    except Exception as e:
+        print(f"❌ Google GenAI 图像生成失败: {str(e)}")
+
+
+def test_google_vertex_ai_image_generation():
+    """测试 Google Vertex AI 图像生成 (对比)"""
+    print("\n🎨 测试 Google Vertex AI 图像生成...")
+
+    try:
+        client = TamarModelClient()
+
+        request = ModelRequest(
+            provider=ProviderType.GOOGLE,
+            channel=Channel.VERTEXAI,
+            invoke_type=InvokeType.IMAGE_GENERATION,
+            model="imagegeneration@006",
+            prompt="一座雄伟的雪山在黄昏时分",
+            user_context=UserContext(
+                user_id="test_user",
+                org_id="test_org",
+                client_type="test_client"
+            )
+        )
+
+        response = client.invoke(request)
+        print(f"✅ Google Vertex AI 图像生成成功")
+        print(f"   响应类型: {type(response)}")
+        if response.content:
+            print(f"   图像内容长度: {len(str(response.content))}")
+        else:
+            print(f"   响应内容: {str(response)[:200]}...")
+
+    except Exception as e:
+        print(f"❌ Google Vertex AI 图像生成失败: {str(e)}")
+
+
 async def test_google_streaming():
     """测试 Google 流式响应"""
     print("\n📡 测试 Google 流式响应...")
@@ -223,6 +288,78 @@ async def test_azure_streaming():
         print(f"❌ Azure 流式响应失败: {str(e)}")
 
 
+async def test_google_genai_image_generation_async():
+    """测试异步 Google GenAI 图像生成"""
+    print("\n🎨 测试异步 Google GenAI 图像生成...")
+
+    try:
+        async with AsyncTamarModelClient() as client:
+            request = ModelRequest(
+                provider=ProviderType.GOOGLE,
+                channel=Channel.AI_STUDIO,
+                invoke_type=InvokeType.IMAGE_GENERATION_GENAI,
+                model="imagen-3.0-generate-002",
+                prompt="现代城市夜景，霓虹灯闪烁，繁华热闹的街道",
+                user_context=UserContext(
+                    user_id="test_user_async",
+                    org_id="test_org",
+                    client_type="test_client_async"
+                )
+            )
+
+            response = await client.invoke(request)
+            print(f"✅ 异步 Google GenAI 图像生成调用成功")
+            print(f"   响应类型: {type(response)}")
+
+            # 对于图像生成，是直接的响应对象，不是流式
+            if response.content:
+                print(f"   图像内容长度: {len(str(response.content))}")
+                print(f"   图像生成成功！")
+            elif response.error:
+                print(f"   错误: {response.error}")
+            else:
+                print(f"   响应内容: {str(response)[:200]}...")
+
+    except Exception as e:
+        print(f"❌ 异步 Google GenAI 图像生成失败: {str(e)}")
+
+
+async def test_google_vertex_ai_image_generation_async():
+    """测试异步 Google Vertex AI 图像生成 (对比)"""
+    print("\n🎨 测试异步 Google Vertex AI 图像生成...")
+
+    try:
+        async with AsyncTamarModelClient() as client:
+            request = ModelRequest(
+                provider=ProviderType.GOOGLE,
+                channel=Channel.VERTEXAI,
+                invoke_type=InvokeType.IMAGE_GENERATION,
+                model="imagegeneration@006",
+                prompt="宁静的湖泊倒映着夕阳，周围环绕着青山绿树",
+                user_context=UserContext(
+                    user_id="test_user_async",
+                    org_id="test_org",
+                    client_type="test_client_async"
+                )
+            )
+
+            response = await client.invoke(request)
+            print(f"✅ 异步 Google Vertex AI 图像生成调用成功")
+            print(f"   响应类型: {type(response)}")
+
+            # 对于图像生成，是直接的响应对象，不是流式
+            if response.content:
+                print(f"   图像内容长度: {len(str(response.content))}")
+                print(f"   图像生成成功！")
+            elif response.error:
+                print(f"   错误: {response.error}")
+            else:
+                print(f"   响应内容: {str(response)[:200]}...")
+
+    except Exception as e:
+        print(f"❌ 异步 Google Vertex AI 图像生成失败: {str(e)}")
+
+
 def test_sync_batch_requests():
     """测试同步批量请求"""
     print("\n📦 测试同步批量请求...")
@@ -332,6 +469,19 @@ async def test_batch_requests():
                         ],
                         custom_id="google-vertex-ai-1",
                     ),
+                    # Google GenAI 图像生成请求
+                    BatchModelRequestItem(
+                        provider=ProviderType.GOOGLE,
+                        channel=Channel.AI_STUDIO,
+                        invoke_type=InvokeType.IMAGE_GENERATION_GENAI,
+                        model="imagen-3.0-generate-001",
+                        prompt="一朵美丽的玫瑰花在阳光下绽放",
+                        config=types.GenerateImagesConfig(
+                            number_of_images=1,
+                            aspect_ratio="1:1"
+                        ),
+                        custom_id="google-genai-image-1",
+                    ),
                     # Azure OpenAI 请求
                     BatchModelRequestItem(
                         provider=ProviderType.AZURE,
@@ -378,6 +528,120 @@ async def test_batch_requests():
         print(f"❌ 批量请求失败: {str(e)}")
 
 
+async def test_image_generation_batch():
+    """测试图像生成批量请求 - 同时测试 GenAI、Vertex AI 图像生成"""
+    print("\n🖼️  测试图像生成批量请求...")
+
+    try:
+        from tamar_model_client.schemas import BatchModelRequest, BatchModelRequestItem
+
+        async with AsyncTamarModelClient() as client:
+            # 构建图像生成批量请求
+            batch_request = BatchModelRequest(
+                user_context=UserContext(
+                    user_id="test_image_batch",
+                    org_id="test_org",
+                    client_type="image_test_client"
+                ),
+                items=[
+                    # Google GenAI 图像生成请求 1
+                    BatchModelRequestItem(
+                        provider=ProviderType.GOOGLE,
+                        channel=Channel.AI_STUDIO,
+                        invoke_type=InvokeType.IMAGE_GENERATION_GENAI,
+                        model="imagen-3.0-generate-001",
+                        prompt="一只可爱的小狗在公园里奔跑",
+                        config=types.GenerateImagesConfig(
+                            number_of_images=1,
+                            aspect_ratio="1:1",
+                            safety_filter_level="block_some"
+                        ),
+                        custom_id="genai-dog-1",
+                    ),
+                    # Google GenAI 图像生成请求 2
+                    BatchModelRequestItem(
+                        provider=ProviderType.GOOGLE,
+                        channel=Channel.AI_STUDIO,
+                        invoke_type=InvokeType.IMAGE_GENERATION_GENAI,
+                        model="imagen-3.0-generate-001",
+                        prompt="美丽的樱花盛开在春天的公园里",
+                        config=types.GenerateImagesConfig(
+                            number_of_images=1,
+                            aspect_ratio="16:9"
+                        ),
+                        custom_id="genai-sakura-1",
+                    ),
+                    # Google Vertex AI 图像生成请求 1
+                    BatchModelRequestItem(
+                        provider=ProviderType.GOOGLE,
+                        channel=Channel.VERTEXAI,
+                        invoke_type=InvokeType.IMAGE_GENERATION,
+                        model="imagegeneration@006",
+                        prompt="壮丽的山峦在夕阳西下时的景色",
+                        number_of_images=1,
+                        aspect_ratio="16:9",
+                        safety_filter_level="block_some",
+                        custom_id="vertex-mountain-1",
+                    ),
+                    # Google Vertex AI 图像生成请求 2
+                    BatchModelRequestItem(
+                        provider=ProviderType.GOOGLE,
+                        channel=Channel.VERTEXAI,
+                        invoke_type=InvokeType.IMAGE_GENERATION,
+                        model="imagegeneration@006",
+                        prompt="宁静的海滩上有椰子树和海浪",
+                        number_of_images=1,
+                        aspect_ratio="1:1",
+                        safety_filter_level="block_some",
+                        custom_id="vertex-beach-1",
+                    )
+                ]
+            )
+
+            # 执行批量图像生成请求
+            print(f"   发送批量图像生成请求 (共{len(batch_request.items)}个)...")
+            batch_response = await client.invoke_batch(batch_request)
+
+            print(f"✅ 批量图像生成请求成功")
+            print(f"   请求数量: {len(batch_request.items)}")
+            print(f"   响应数量: {len(batch_response.responses)}")
+            print(f"   批量请求ID: {batch_response.request_id}")
+
+            # 详细显示每个图像生成结果
+            genai_success = 0
+            vertex_success = 0
+            total_errors = 0
+            
+            for i, response in enumerate(batch_response.responses):
+                print(f"\n   图像生成 {i + 1}:")
+                print(f"   - custom_id: {response.custom_id}")
+                print(f"   - 有错误: {'是' if response.error else '否'}")
+                
+                if response.error:
+                    total_errors += 1
+                    print(f"   - 错误信息: {response.error}")
+                else:
+                    if response.content:
+                        print(f"   - 图像内容长度: {len(str(response.content))}")
+                        print(f"   - ✅ 图像生成成功！")
+                        
+                        # 统计不同类型的成功数
+                        if "genai" in response.custom_id:
+                            genai_success += 1
+                        elif "vertex" in response.custom_id:
+                            vertex_success += 1
+                    else:
+                        print(f"   - 响应预览: {str(response)[:100]}...")
+
+            print(f"\n📊 图像生成批量测试统计:")
+            print(f"   - GenAI 图像生成成功: {genai_success}/2")
+            print(f"   - Vertex AI 图像生成成功: {vertex_success}/2")
+            print(f"   - 总错误数: {total_errors}")
+
+    except Exception as e:
+        print(f"❌ 批量图像生成请求失败: {str(e)}")
+
+
 def test_concurrent_requests(num_requests: int = 150):
     """测试并发请求
     
@@ -513,7 +777,7 @@ async def test_async_concurrent_requests(num_requests: int = 150):
     failed_requests = 0
     request_times: List[float] = []
     errors: Dict[str, int] = {}
-    trace_id = "8885588866668888886666888888866666668888"
+    trace_id = "88888888888888888333333888888883333388888"
 
     # 异步锁
     stats_lock = asyncio.Lock()
@@ -542,7 +806,7 @@ async def test_async_concurrent_requests(num_requests: int = 150):
                 config={"temperature": 0.1}
             )
 
-            response = await client.invoke(request, timeout=300000.0)
+            response = await client.invoke(request, timeout=300000.0, request_id=trace_id)
             duration = time.time() - start_time
             return (True, duration, "")
 
@@ -635,26 +899,26 @@ async def test_async_batch_with_circuit_breaker_v2(num_requests: int = 10):
         num_requests: 要发送的请求数，默认10个
     """
     print(f"\n🔥 测试熔断器功能 - 改进版 ({num_requests} 个独立请求)...")
-
+    
     # 保存原始环境变量
     import os
     original_env = {}
-    env_vars = ['MODEL_CLIENT_RESILIENT_ENABLED', 'MODEL_CLIENT_HTTP_FALLBACK_URL',
+    env_vars = ['MODEL_CLIENT_RESILIENT_ENABLED', 'MODEL_CLIENT_HTTP_FALLBACK_URL', 
                 'MODEL_CLIENT_CIRCUIT_BREAKER_THRESHOLD', 'MODEL_CLIENT_CIRCUIT_BREAKER_TIMEOUT']
     for var in env_vars:
         original_env[var] = os.environ.get(var)
-
+    
     # 设置环境变量以启用熔断器和HTTP fallback
     os.environ['MODEL_CLIENT_RESILIENT_ENABLED'] = 'true'
     os.environ['MODEL_CLIENT_HTTP_FALLBACK_URL'] = 'http://localhost:8000'
     os.environ['MODEL_CLIENT_CIRCUIT_BREAKER_THRESHOLD'] = '3'  # 3次失败后触发熔断
     os.environ['MODEL_CLIENT_CIRCUIT_BREAKER_TIMEOUT'] = '30'  # 熔断器30秒后恢复
-
+    
     print(f"   环境变量设置:")
     print(f"   - MODEL_CLIENT_RESILIENT_ENABLED: {os.environ.get('MODEL_CLIENT_RESILIENT_ENABLED')}")
     print(f"   - MODEL_CLIENT_HTTP_FALLBACK_URL: {os.environ.get('MODEL_CLIENT_HTTP_FALLBACK_URL')}")
-    print(f"   - 熔断阈值: 1 次失败")
-
+    print(f"   - 熔断阈值: 3 次失败")
+    
     # 统计变量
     total_requests = 0
     successful_requests = 0
@@ -663,24 +927,25 @@ async def test_async_batch_with_circuit_breaker_v2(num_requests: int = 10):
     http_fallback_used = 0
     request_times: List[float] = []
     errors: Dict[str, int] = {}
-
+    
     try:
         # 创建一个共享的异步客户端（启用熔断器）
         async with AsyncTamarModelClient() as client:
             print(f"\n   熔断器配置:")
             print(f"   - 启用状态: {getattr(client, 'resilient_enabled', False)}")
             print(f"   - HTTP Fallback URL: {getattr(client, 'http_fallback_url', 'None')}")
-
+            
             for i in range(num_requests):
                 start_time = time.time()
-
+                
                 try:
                     # 前4个请求使用错误的model来触发失败
-                    if i < 6:
+                    if i < 4:
                         request = ModelRequest(
-                            provider=ProviderType.OPENAI,
+                            provider=ProviderType.GOOGLE,
+                            invoke_type=InvokeType.GENERATION,
                             model="invalid-model-to-trigger-error",  # 无效模型
-                            input=f"测试失败请求 {i + 1}",
+                            contents=f"测试失败请求 {i+1}",
                             user_context=UserContext(
                                 user_id=f"circuit_test_{i}",
                                 org_id="test_org_circuit",
@@ -691,8 +956,9 @@ async def test_async_batch_with_circuit_breaker_v2(num_requests: int = 10):
                         # 后续请求使用正确的model
                         request = ModelRequest(
                             provider=ProviderType.GOOGLE,
+                            invoke_type=InvokeType.GENERATION,
                             model="tamar-google-gemini-flash-lite",
-                            contents=f"测试请求 {i + 1}: 计算 {i} + {i}",
+                            contents=f"测试请求 {i+1}: 计算 {i} + {i}",
                             user_context=UserContext(
                                 user_id=f"circuit_test_{i}",
                                 org_id="test_org_circuit",
@@ -700,17 +966,17 @@ async def test_async_batch_with_circuit_breaker_v2(num_requests: int = 10):
                             ),
                             config={"temperature": 0.1}
                         )
-
-                    print(f"\n   📤 发送请求 {i + 1}/{num_requests}...")
+                    
+                    print(f"\n   📤 发送请求 {i+1}/{num_requests}...")
                     response = await client.invoke(request, timeout=10000)
-
+                    
                     duration = time.time() - start_time
                     request_times.append(duration)
                     total_requests += 1
                     successful_requests += 1
-
-                    print(f"   ✅ 请求 {i + 1} 成功 - 耗时: {duration:.2f}秒")
-
+                    
+                    print(f"   ✅ 请求 {i+1} 成功 - 耗时: {duration:.2f}秒")
+                    
                     # 检查是否通过HTTP fallback
                     if hasattr(client, 'resilient_enabled') and client.resilient_enabled:
                         try:
@@ -720,19 +986,19 @@ async def test_async_batch_with_circuit_breaker_v2(num_requests: int = 10):
                                 print(f"      (通过HTTP fallback)")
                         except:
                             pass
-
+                    
                 except Exception as e:
                     duration = time.time() - start_time
                     request_times.append(duration)
                     total_requests += 1
                     failed_requests += 1
-
+                    
                     error_type = type(e).__name__
                     errors[error_type] = errors.get(error_type, 0) + 1
-
-                    print(f"   ❌ 请求 {i + 1} 失败: {error_type} - {str(e)[:100]}")
+                    
+                    print(f"   ❌ 请求 {i+1} 失败: {error_type} - {str(e)[:100]}")
                     print(f"      耗时: {duration:.2f}秒")
-
+                
                 # 检查熔断器状态
                 if hasattr(client, 'resilient_enabled') and client.resilient_enabled:
                     try:
@@ -740,28 +1006,28 @@ async def test_async_batch_with_circuit_breaker_v2(num_requests: int = 10):
                         if metrics and 'circuit_breaker' in metrics:
                             state = metrics['circuit_breaker']['state']
                             failures = metrics['circuit_breaker']['failure_count']
-
+                            
                             if state == 'open' and not circuit_breaker_opened:
                                 circuit_breaker_opened = True
                                 print(f"   🔻 熔断器已打开！失败次数: {failures}")
-
+                            
                             print(f"      熔断器: {state}, 失败计数: {failures}")
                     except Exception as e:
                         print(f"      获取熔断器状态失败: {e}")
-
+                
                 # 请求之间短暂等待
                 await asyncio.sleep(0.2)
-
+            
             # 最终统计
             print(f"\n📊 熔断器测试结果:")
             print(f"   总请求数: {total_requests}")
             print(f"   成功请求: {successful_requests}")
             print(f"   失败请求: {failed_requests}")
-
+            
             print(f"\n   🔥 熔断器统计:")
             print(f"   - 熔断器是否触发: {'是' if circuit_breaker_opened else '否'}")
             print(f"   - HTTP fallback使用次数: {http_fallback_used}")
-
+            
             # 获取最终状态
             if hasattr(client, 'resilient_enabled') and client.resilient_enabled:
                 try:
@@ -771,17 +1037,17 @@ async def test_async_batch_with_circuit_breaker_v2(num_requests: int = 10):
                         print(f"   - 总失败次数: {final_metrics['circuit_breaker']['failure_count']}")
                 except Exception as e:
                     print(f"   - 获取最终状态失败: {e}")
-
+            
             if errors:
                 print(f"\n   错误统计:")
                 for error_type, count in sorted(errors.items(), key=lambda x: x[1], reverse=True):
                     print(f"   - {error_type}: {count} 次")
-
+    
     except Exception as e:
         print(f"❌ 测试失败: {str(e)}")
         import traceback
         traceback.print_exc()
-
+    
     finally:
         # 恢复原始环境变量
         for var, value in original_env.items():
@@ -1205,34 +1471,36 @@ async def main():
 
     try:
         # 同步测试
-        test_google_ai_studio()
-        test_google_vertex_ai()
-        test_azure_openai()
-
+        # test_google_ai_studio()
+        # test_google_vertex_ai()
+        # test_azure_openai()
+        
+        # 新增：图像生成测试
+        # test_google_genai_image_generation()
+        # test_google_vertex_ai_image_generation()
+        
         # 同步批量测试
-        test_sync_batch_requests()
+        # test_sync_batch_requests()
 
         # 异步流式测试
-        await asyncio.wait_for(test_google_streaming(), timeout=60.0)
-        await asyncio.wait_for(test_azure_streaming(), timeout=60.0)
-
-        # 异步批量测试
-        await asyncio.wait_for(test_batch_requests(), timeout=120.0)
+        # await asyncio.wait_for(test_google_streaming(), timeout=60.0)
+        # await asyncio.wait_for(test_azure_streaming(), timeout=60.0)
+        
+        # 新增：异步图像生成测试
+        # await asyncio.wait_for(test_google_genai_image_generation_async(), timeout=120.0)
+        # await asyncio.wait_for(test_google_vertex_ai_image_generation_async(), timeout=120.0)
+        #
+        # # 异步批量测试
+        # await asyncio.wait_for(test_batch_requests(), timeout=120.0)
+        #
+        # # 新增：图像生成批量测试
+        # await asyncio.wait_for(test_image_generation_batch(), timeout=180.0)
 
         # 同步并发测试
-        test_concurrent_requests(2)  # 测试150个并发请求
-
+        # test_concurrent_requests(2)  # 测试150个并发请求
+        #
         # # 异步并发测试
-        await test_async_concurrent_requests(2)  # 测试50个异步并发请求（复用连接）
-
-        # 异步并发测试 - 独立客户端模式
-        # await test_async_concurrent_requests_independent_clients(30)  # 测试30个独立客户端并发请求
-
-        # 异步批量测试 - 熔断器模式（原版）
-        # await test_async_batch_with_circuit_breaker(10, 5)  # 测试5个批量请求，每批10个请求
-
-        # 熔断器测试 - 改进版（使用单个请求）
-        # await test_async_batch_with_circuit_breaker_v2(10)  # 测试10个独立请求触发熔断
+        # await test_async_concurrent_requests(2)  # 测试50个异步并发请求（复用连接）
 
         print("\n✅ 测试完成")