tamar-model-client 0.2.5__tar.gz → 0.2.7__tar.gz

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tamar-model-client
-Version: 0.2.5
+Version: 0.2.7
 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
@@ -16,8 +16,8 @@ Requires-Dist: grpcio-tools~=1.67.1
 Requires-Dist: pydantic
 Requires-Dist: PyJWT
 Requires-Dist: nest_asyncio
-Requires-Dist: openai
-Requires-Dist: google-genai
+Requires-Dist: openai>=1.99.3
+Requires-Dist: google-genai>=1.29.0
 Requires-Dist: requests>=2.25.0
 Requires-Dist: aiohttp>=3.7.0
 Dynamic: author

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="tamar-model-client",
-    version="0.2.5",
+    version="0.2.7",
     description="A Python SDK for interacting with the Model Manager gRPC service",
     author="Oscar Ou",
     author_email="oscar.ou@tamaredge.ai",
@@ -17,8 +17,8 @@ setup(
         "pydantic",
         "PyJWT",
         "nest_asyncio",
-        "openai",
-        "google-genai",
+        "openai>=1.99.3",
+        "google-genai>=1.29.0",
         "requests>=2.25.0",  # HTTP降级功能（同步）
         "aiohttp>=3.7.0",    # HTTP降级功能（异步）
     ],

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client/__init__.py

@@ -3,6 +3,7 @@ from .async_client import AsyncTamarModelClient
 from .exceptions import ModelManagerClientError, ConnectionError, ValidationError
 from .json_formatter import JSONFormatter
 from . import logging_icons
+from .tool_call_helper import ToolCallHelper
 
 __all__ = [
     "TamarModelClient",
@@ -12,4 +13,5 @@ __all__ = [
     "ValidationError",
     "JSONFormatter",
     "logging_icons",
+    "ToolCallHelper",
 ]

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client/async_client.py

@@ -623,7 +623,7 @@ class AsyncTamarModelClient(BaseClient, AsyncHttpFallbackMixin):
             origin_request_id
         )
         stream_iter = self.stub.Invoke(request, metadata=fresh_metadata, timeout=invoke_timeout).__aiter__()
-        chunk_timeout = 30.0  # 单个数据块的超时时间
+        chunk_timeout = self.stream_chunk_timeout  # 单个数据块的超时时间
         
         try:
             while True:

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client/core/base_client.py

@@ -296,6 +296,10 @@ class BaseClient(ABC):
                 if hasattr(grpc.StatusCode, error_name):
                     self.never_fallback_errors.add(getattr(grpc.StatusCode, error_name))
 
+        # 流式响应单个数据块的超时时间（秒）
+        # AI模型生成可能需要更长时间，默认设置为120秒
+        self.stream_chunk_timeout = float(os.getenv('MODEL_CLIENT_STREAM_CHUNK_TIMEOUT', '120.0'))
+
         if self.fast_fallback_enabled:
             self.logger.info(
                 "🚀 Fast fallback enabled",

tamar_model_client-0.2.7/tamar_model_client/core/response_handler.py

@@ -0,0 +1,219 @@
+"""
+Response handling logic for Tamar Model Client
+
+This module provides utilities for processing gRPC responses and
+converting them to client response objects.
+"""
+
+import json
+from typing import Optional, Dict, Any
+
+from ..schemas import ModelResponse, BatchModelResponse
+
+
+class ResponseHandler:
+    """
+    响应处理器
+    
+    负责将 gRPC 响应转换为客户端响应对象，
+    包括 JSON 解析、错误处理和数据结构转换。
+    """
+    
+    @staticmethod
+    def build_model_response(grpc_response) -> ModelResponse:
+        """
+        从 gRPC 响应构建增强的 ModelResponse 对象
+        
+        新增功能：
+        1. 自动提取 tool_calls（对标 OpenAI SDK）
+        2. 提取 finish_reason（对标 OpenAI SDK）
+        3. 支持多种 provider 格式转换
+        
+        Args:
+            grpc_response: gRPC 服务返回的响应对象
+            
+        Returns:
+            ModelResponse: 增强的客户端响应对象
+        """
+        raw_response = ResponseHandler._parse_json_field(grpc_response.raw_response)
+        
+        # 提取 tool_calls 和 finish_reason
+        tool_calls = None
+        finish_reason = None
+        
+        if raw_response and isinstance(raw_response, dict):
+            # OpenAI/Azure OpenAI 格式
+            if 'choices' in raw_response and raw_response['choices']:
+                choice = raw_response['choices'][0]
+                
+                # 提取 tool_calls
+                if 'message' in choice and 'tool_calls' in choice['message']:
+                    tool_calls = choice['message']['tool_calls']
+                
+                # 提取 finish_reason
+                if 'finish_reason' in choice:
+                    finish_reason = choice['finish_reason']
+            
+            # Google AI 格式适配
+            elif 'candidates' in raw_response and raw_response['candidates']:
+                candidate = raw_response['candidates'][0]
+                
+                # Google 格式的 function calls 映射
+                if 'content' in candidate and 'parts' in candidate['content']:
+                    parts = candidate['content']['parts']
+                    google_tool_calls = []
+                    
+                    for i, part in enumerate(parts):
+                        if 'functionCall' in part:
+                            # 转换为 OpenAI 兼容格式
+                            function_call = part['functionCall']
+                            google_tool_calls.append({
+                                'id': f"call_{i}_{function_call.get('name', 'unknown')}",
+                                'type': 'function',
+                                'function': {
+                                    'name': function_call.get('name', ''),
+                                    'arguments': json.dumps(function_call.get('args', {}))
+                                }
+                            })
+                    
+                    if google_tool_calls:
+                        tool_calls = google_tool_calls
+                
+                # Google 的 finish_reason
+                if 'finishReason' in candidate:
+                    # 映射 Google 格式到标准格式
+                    google_reason = candidate['finishReason']
+                    finish_reason_mapping = {
+                        'STOP': 'stop',
+                        'MAX_TOKENS': 'length',
+                        'SAFETY': 'content_filter',
+                        'RECITATION': 'content_filter'
+                    }
+                    finish_reason = finish_reason_mapping.get(google_reason, google_reason.lower())
+                    
+                    # 如果有工具调用，设置 finish_reason 为 tool_calls
+                    if tool_calls:
+                        finish_reason = 'tool_calls'
+        
+        return ModelResponse(
+            content=grpc_response.content,
+            usage=ResponseHandler._parse_json_field(grpc_response.usage),
+            error=grpc_response.error or None,
+            raw_response=raw_response,
+            request_id=grpc_response.request_id if grpc_response.request_id else None,
+            tool_calls=tool_calls,
+            finish_reason=finish_reason
+        )
+    
+    @staticmethod
+    def build_batch_response(grpc_response) -> BatchModelResponse:
+        """
+        从 gRPC 批量响应构建 BatchModelResponse 对象
+        
+        Args:
+            grpc_response: gRPC 服务返回的批量响应对象
+            
+        Returns:
+            BatchModelResponse: 客户端批量响应对象
+        """
+        responses = []
+        for response_item in grpc_response.items:
+            model_response = ResponseHandler.build_model_response(response_item)
+            responses.append(model_response)
+        
+        return BatchModelResponse(
+            responses=responses,
+            request_id=grpc_response.request_id if grpc_response.request_id else None
+        )
+    
+    @staticmethod
+    def _parse_json_field(json_str: Optional[str]) -> Optional[Dict[str, Any]]:
+        """
+        安全地解析 JSON 字符串
+        
+        Args:
+            json_str: 待解析的 JSON 字符串
+            
+        Returns:
+            Optional[Dict[str, Any]]: 解析后的字典，或 None（如果输入为空）
+        """
+        if not json_str:
+            return None
+        
+        try:
+            return json.loads(json_str)
+        except json.JSONDecodeError:
+            # 如果解析失败，返回原始字符串作为错误信息
+            return {"error": "JSON parse error", "raw": json_str}
+    
+    @staticmethod
+    def build_log_data(
+        model_request,
+        response: Optional[ModelResponse] = None,
+        duration: Optional[float] = None,
+        error: Optional[Exception] = None,
+        stream_stats: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        构建日志数据
+        
+        为请求和响应日志构建结构化的数据字典。
+        
+        Args:
+            model_request: 原始请求对象
+            response: 响应对象（可选）
+            duration: 请求持续时间（秒）
+            error: 错误对象（可选）
+            stream_stats: 流式响应统计信息（可选）
+            
+        Returns:
+            Dict[str, Any]: 日志数据字典
+        """
+        data = {
+            "provider": model_request.provider.value,
+            "invoke_type": model_request.invoke_type.value,
+            "model": getattr(model_request, 'model', None),
+            "stream": getattr(model_request, 'stream', False),
+        }
+        
+        # 添加用户上下文信息（如果有）
+        if hasattr(model_request, 'user_context'):
+            data.update({
+                "org_id": model_request.user_context.org_id,
+                "user_id": model_request.user_context.user_id,
+                "client_type": model_request.user_context.client_type
+            })
+        
+        # 添加请求中的 tool 信息
+        if hasattr(model_request, 'tools') and model_request.tools:
+            data["tools_count"] = len(model_request.tools) if isinstance(model_request.tools, list) else 1
+            data["has_tools"] = True
+        
+        if hasattr(model_request, 'tool_choice') and model_request.tool_choice:
+            data["tool_choice"] = str(model_request.tool_choice)
+        
+        # 添加响应信息
+        if response:
+            if hasattr(response, 'content') and response.content:
+                data["content_length"] = len(response.content)
+            if hasattr(response, 'usage'):
+                data["usage"] = response.usage
+            
+            # 新增：tool_calls 相关日志
+            if hasattr(response, 'tool_calls') and response.tool_calls:
+                data["tool_calls_count"] = len(response.tool_calls)
+                data["has_tool_calls"] = True
+                
+            if hasattr(response, 'finish_reason') and response.finish_reason:
+                data["finish_reason"] = response.finish_reason
+        
+        # 添加流式响应统计
+        if stream_stats:
+            data.update(stream_stats)
+        
+        # 添加错误信息
+        if error:
+            data["error_type"] = type(error).__name__
+            data["error_message"] = str(error)
+        
+        return data

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client/schemas/__init__.py

@@ -2,11 +2,12 @@
 Schema definitions for the API
 """
 
-from .inputs import UserContext, ModelRequest, BatchModelRequestItem, BatchModelRequest
+from .inputs import UserContext, ModelRequest, BatchModelRequestItem, BatchModelRequest, TamarFileIdInput
 from .outputs import ModelResponse, BatchModelResponse
 
 __all__ = [
     # Model Inputs
+    "TamarFileIdInput",
     "UserContext",
     "ModelRequest",
     "BatchModelRequestItem",

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client/schemas/inputs.py

@@ -26,6 +26,10 @@ class UserContext(BaseModel):
     client_type: str  # 客户端类型，这里记录的是哪个服务请求过来的
 
 
+class TamarFileIdInput(BaseModel):
+    file_id: str
+
+
 class GoogleGenAiInput(BaseModel):
     model: str
     contents: Union[types.ContentListUnion, types.ContentListUnionDict]
@@ -179,11 +183,11 @@ class OpenAIImagesInput(BaseModel):
 
 
 class OpenAIImagesEditInput(BaseModel):
-    image: Union[FileTypes, List[FileTypes]]
+    image: Union[FileTypes, List[FileTypes], TamarFileIdInput, List[TamarFileIdInput]]
     prompt: str
     background: Optional[Literal["transparent", "opaque", "auto"]] | NotGiven = NOT_GIVEN
     input_fidelity: Optional[Literal["high", "low"]] | NotGiven = NOT_GIVEN
-    mask: FileTypes | NotGiven = NOT_GIVEN
+    mask: FileTypes | TamarFileIdInput | NotGiven = NOT_GIVEN
     model: Union[str, ImageModel, None] | NotGiven = NOT_GIVEN
     n: Optional[int] | NotGiven = NOT_GIVEN
     output_compression: Optional[int] | NotGiven = NOT_GIVEN
@@ -277,7 +281,7 @@ class ModelRequestInput(BaseRequest):
     config: Optional[Union[types.GenerateContentConfigOrDict, types.GenerateImagesConfigOrDict]] = None
 
     # Images（OpenAI Images / Images Edit / Google Vertex Images 合并）
-    image: Optional[Union[FileTypes, List[FileTypes]]] = None
+    image: Optional[Union[FileTypes, List[FileTypes], TamarFileIdInput, List[TamarFileIdInput]]] = None
     # background 同名字段合并：Responses 的 bool（后台任务）+ Images 的透明度枚举
     background: Optional[Union[bool, Literal["transparent", "opaque", "auto"], NotGiven]] = NOT_GIVEN
     moderation: Optional[Union[Literal["low", "auto"], NotGiven]] = NOT_GIVEN
@@ -285,7 +289,7 @@ class ModelRequestInput(BaseRequest):
     output_compression: Optional[Union[int, NotGiven]] = NOT_GIVEN
     output_format: Optional[Union[Literal["png", "jpeg", "webp"], NotGiven]] = NOT_GIVEN
     partial_images: Optional[Union[int, NotGiven]] = NOT_GIVEN
-    mask: Union[FileTypes, NotGiven] = NOT_GIVEN
+    mask: Union[FileTypes, TamarFileIdInput, NotGiven] = NOT_GIVEN
     negative_prompt: Optional[str] = None
     aspect_ratio: Optional[Literal["1:1", "9:16", "16:9", "4:3", "3:4"]] = None
     guidance_scale: Optional[float] = None

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client/schemas/outputs.py

@@ -15,8 +15,27 @@ class BaseResponse(BaseModel):
 
 
 class ModelResponse(BaseResponse):
+    """增强的模型响应类，对标 OpenAI SDK 的 Tool Call 支持"""
+    
     model_config = ConfigDict(arbitrary_types_allowed=True)
+    
     request_id: Optional[str] = None  # 请求ID，用于跟踪请求
+    
+    # 新增字段 - 对标 OpenAI SDK
+    tool_calls: Optional[List[Dict[str, Any]]] = None
+    """Tool calls 列表，对应 OpenAI SDK 的 message.tool_calls"""
+    
+    finish_reason: Optional[str] = None
+    """完成原因，对应 OpenAI SDK 的 choice.finish_reason"""
+    
+    # 基础便利方法
+    def has_tool_calls(self) -> bool:
+        """检查响应是否包含 tool calls
+        
+        Returns:
+            bool: 如果包含 tool calls 返回 True
+        """
+        return bool(self.tool_calls and len(self.tool_calls) > 0)
 
 
 class BatchModelResponse(BaseModel):

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client/sync_client.py

@@ -786,7 +786,7 @@ class TamarModelClient(BaseClient, HttpFallbackMixin):
         fetch_thread.daemon = True
         fetch_thread.start()
         
-        chunk_timeout = 30.0  # 单个数据块的超时时间
+        chunk_timeout = self.stream_chunk_timeout  # 单个数据块的超时时间
         
         while True:
             # 检查是否有异常

tamar_model_client-0.2.7/tamar_model_client/tool_call_helper.py

@@ -0,0 +1,169 @@
+"""
+Tool Call 实用工具类
+
+提供简化 Tool Call 使用的基础工具方法，减少常见错误和样板代码。
+注意：本工具类仅提供数据处理便利，不包含自动执行功能。
+"""
+
+import json
+from typing import List, Dict, Any, Optional
+
+from .schemas.outputs import ModelResponse
+
+
+class ToolCallHelper:
+    """Tool Call 实用工具类
+    
+    提供基础的数据处理方法，对标 OpenAI SDK 的使用体验。
+    """
+    
+    @staticmethod
+    def create_function_tool(
+        name: str, 
+        description: str, 
+        parameters: Dict[str, Any],
+        strict: Optional[bool] = None
+    ) -> Dict[str, Any]:
+        """创建函数工具定义（对标 OpenAI SDK 的工具定义格式）
+        
+        Args:
+            name: 函数名称
+            description: 函数描述
+            parameters: 函数参数的 JSON Schema
+            strict: 是否启用严格模式（OpenAI Structured Outputs）
+            
+        Returns:
+            ChatCompletionToolParam: 工具定义对象
+            
+        Example:
+            >>> weather_tool = ToolCallHelper.create_function_tool(
+            ...     name="get_weather",
+            ...     description="获取指定城市的天气信息",
+            ...     parameters={
+            ...         "type": "object",
+            ...         "properties": {
+            ...             "location": {"type": "string", "description": "城市名称"}
+            ...         },
+            ...         "required": ["location"]
+            ...     }
+            ... )
+        """
+        tool_def = {
+            "type": "function",
+            "function": {
+                "name": name,
+                "description": description,
+                "parameters": parameters
+            }
+        }
+        
+        if strict is not None:
+            tool_def["function"]["strict"] = strict
+            
+        return tool_def
+    
+    @staticmethod
+    def create_tool_response_message(
+        tool_call_id: str, 
+        content: str,
+        name: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """创建工具响应消息（对标 OpenAI SDK 的消息格式）
+        
+        Args:
+            tool_call_id: 工具调用 ID
+            content: 工具执行结果
+            name: 工具名称（可选）
+            
+        Returns:
+            ChatCompletionMessageParam: 工具响应消息
+            
+        Example:
+            >>> tool_message = ToolCallHelper.create_tool_response_message(
+            ...     tool_call_id="call_123",
+            ...     content="北京今天晴天，25°C",
+            ...     name="get_weather"
+            ... )
+        """
+        message = {
+            "role": "tool",
+            "tool_call_id": tool_call_id,
+            "content": content
+        }
+        
+        if name:
+            message["name"] = name
+            
+        return message
+    
+    @staticmethod
+    def parse_function_arguments(tool_call: Dict[str, Any]) -> Dict[str, Any]:
+        """安全解析函数参数（解决 OpenAI SDK 需要手动 json.loads 的痛点）
+        
+        Args:
+            tool_call: 工具调用对象
+            
+        Returns:
+            Dict[str, Any]: 解析后的参数字典
+            
+        Raises:
+            ValueError: 不支持的工具类型或参数解析失败
+            
+        Example:
+            >>> tool_call = response.tool_calls[0]
+            >>> arguments = ToolCallHelper.parse_function_arguments(tool_call)
+            >>> print(arguments["location"])  # "北京"
+        """
+        if tool_call.get("type") != "function":
+            raise ValueError(f"不支持的工具类型: {tool_call.get('type')}")
+        
+        function = tool_call.get("function", {})
+        arguments_str = function.get("arguments", "{}")
+        
+        try:
+            return json.loads(arguments_str)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"解析工具参数失败: {arguments_str}") from e
+    
+    @staticmethod  
+    def build_messages_with_tool_response(
+        original_messages: List[Dict[str, Any]],
+        assistant_message: ModelResponse, 
+        tool_responses: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """构建包含工具响应的消息列表（简化版工具方法）
+        
+        Args:
+            original_messages: 原始消息列表
+            assistant_message: 包含 tool calls 的助手响应
+            tool_responses: 工具响应列表
+            
+        Returns:
+            List[Dict[str, Any]]: 新的消息列表
+            
+        Example:
+            >>> new_messages = ToolCallHelper.build_messages_with_tool_response(
+            ...     original_messages=request.messages,
+            ...     assistant_message=response,
+            ...     tool_responses=[tool_message]
+            ... )
+            >>> # 然后开发者手动创建新请求发送
+        """
+        new_messages = list(original_messages)
+        
+        # 添加助手的响应消息
+        assistant_msg = {
+            "role": "assistant",
+            "content": assistant_message.content or ""
+        }
+        
+        # 如果有 tool calls，添加到消息中
+        if assistant_message.has_tool_calls():
+            assistant_msg["tool_calls"] = assistant_message.tool_calls
+        
+        new_messages.append(assistant_msg)
+        
+        # 添加工具响应消息
+        new_messages.extend(tool_responses)
+        
+        return new_messages

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tamar-model-client
-Version: 0.2.5
+Version: 0.2.7
 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
@@ -16,8 +16,8 @@ Requires-Dist: grpcio-tools~=1.67.1
 Requires-Dist: pydantic
 Requires-Dist: PyJWT
 Requires-Dist: nest_asyncio
-Requires-Dist: openai
-Requires-Dist: google-genai
+Requires-Dist: openai>=1.99.3
+Requires-Dist: google-genai>=1.29.0
 Requires-Dist: requests>=2.25.0
 Requires-Dist: aiohttp>=3.7.0
 Dynamic: author

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client.egg-info/SOURCES.txt

@@ -9,6 +9,7 @@ tamar_model_client/exceptions.py
 tamar_model_client/json_formatter.py
 tamar_model_client/logging_icons.py
 tamar_model_client/sync_client.py
+tamar_model_client/tool_call_helper.py
 tamar_model_client/utils.py
 tamar_model_client.egg-info/PKG-INFO
 tamar_model_client.egg-info/SOURCES.txt
@@ -35,4 +36,5 @@ tamar_model_client/schemas/inputs.py
 tamar_model_client/schemas/outputs.py
 tests/__init__.py
 tests/test_circuit_breaker.py
-tests/test_google_azure_final.py
+tests/test_google_azure_final.py
+tests/test_tool_call_enhancement.py

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tamar_model_client.egg-info/requires.txt

@@ -3,7 +3,7 @@ grpcio-tools~=1.67.1
 pydantic
 PyJWT
 nest_asyncio
-openai
-google-genai
+openai>=1.99.3
+google-genai>=1.29.0
 requests>=2.25.0
 aiohttp>=3.7.0

{tamar_model_client-0.2.5 → tamar_model_client-0.2.7}/tests/test_google_azure_final.py

@@ -1471,36 +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_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_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)
+        #：异步图像生成测试
+        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个并发请求
-        #
-        # # 异步并发测试
-        # await test_async_concurrent_requests(2)  # 测试50个异步并发请求（复用连接）
+        test_concurrent_requests(2)  # 测试150个并发请求
+
+        # 异步并发测试
+        await test_async_concurrent_requests(2)  # 测试50个异步并发请求（复用连接）
 
         print("\n✅ 测试完成")