如何为Vertex AI SDK for Python调用Gemini生成内容设置超时

目前Vertex AI Python SDK的generate_content方法没有直接暴露超时参数，但可以通过两种方式实现超时控制：

方法一：用Python标准库实现超时包装 #

利用线程池的超时机制，在指定时间内未获得响应则抛出异常，示例代码如下：

from vertexai import init
from vertexai.preview.generative_models import GenerativeModel
from google.oauth2.service_account import Credentials
from concurrent.futures import ThreadPoolExecutor, TimeoutError

# 初始化认证与环境
credentials = Credentials.from_service_account_file(
    'path/to/json/credential/file.json')
init(
    project=credentials.project_id, location='northamerica-northeast1',
    credentials=credentials)
model = GenerativeModel('gemini-pro')

def generate_with_timeout(prompt, timeout_sec=30):
    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(model.generate_content, prompt)
        try:
            return future.result(timeout=timeout_sec)
        except TimeoutError:
            raise Exception(f"生成内容超时，已超过{timeout_sec}秒")

try:
    response = generate_with_timeout("Pick a number", timeout_sec=30)
    print("Pick a number:", response.candidates[0].content.text)
except Exception as e:
    print(f"错误: {str(e)}")

方法二：配置底层gRPC客户端超时 #

Vertex AI SDK基于gRPC通信，可通过设置gRPC的超时配置实现全局或单次调用的超时控制：

from vertexai import init
from vertexai.preview.generative_models import GenerativeModel
from google.oauth2.service_account import Credentials
from google.api_core import grpc_helpers, retry

# 初始化认证与自定义gRPC通道
credentials = Credentials.from_service_account_file(
    'path/to/json/credential/file.json')

# 创建带超时的gRPC通道
channel = grpc_helpers.create_channel(
    "northamerica-northeast1-aiplatform.googleapis.com:443",
    credentials=credentials,
)
channel = grpc_helpers.add_timeout_metadata(channel, timeout=30)

init(
    project=credentials.project_id, location='northamerica-northeast1',
    credentials=credentials,
    channel=channel
)

model = GenerativeModel('gemini-pro')
try:
    # 调用时也可通过retry参数单独设置超时
    response = model.generate_content(
        "Pick a number",
        retry=retry.Retry(deadline=30)
    )
    print("Pick a number:", response.candidates[0].content.text)
except Exception as e:
    print(f"错误: {str(e)}")

注意事项 #

• 方法一的线程包装更灵活，适合单次调用的独立超时控制；
• 方法二的gRPC配置属于全局级别，会影响所有基于该通道的API调用；
• 超时时间需根据实际场景调整，避免过短导致正常请求被中断。

内容的提问来源于stack exchange，提问作者CookieMan