from google import genai from google.genai import types import os GEMINI_API_KEY = os.environ["GEMINI_API_KEY"] GEMINI_BASE_URL = os.environ.get("GEMINI_BASE_URL", "https://generativelanguage.googleapis.com/v1beta") MODEL_ID = os.environ.get("GEMINI_MODEL", "gemini-3.1-pro-preview") def _make_client() -> genai.Client: return genai.Client( api_key=GEMINI_API_KEY, http_options=types.HttpOptions(base_url=GEMINI_BASE_URL), ) def list_models(): client = _make_client() models = client.models.list() model_ids = sorted([m.name for m in models]) print(f"共找到 {len(model_ids)} 个模型:") for model_id in model_ids: print(f" - {model_id}") def test_model( stream: bool = False, thinking: bool = True, prompt: str = "Hello, how are you?", ): client = _make_client() thinking_config = ( types.ThinkingConfig(include_thoughts=True, thinking_budget=5000) if thinking else types.ThinkingConfig(include_thoughts=False) ) config = types.GenerateContentConfig(thinking_config=thinking_config) if stream: in_thinking = False for chunk in client.models.generate_content_stream( model=MODEL_ID, contents=prompt, config=config ): if not chunk.candidates: continue for part in chunk.candidates[0].content.parts or []: if part.thought: if not in_thinking: print("", flush=True) in_thinking = True print(part.text, end="", flush=True) else: if in_thinking: print("\n\n", flush=True) in_thinking = False print(part.text or "", end="", flush=True) if in_thinking: print("\n", flush=True) print() else: response = client.models.generate_content( model=MODEL_ID, contents=prompt, config=config ) for part in response.candidates[0].content.parts: if part.thought: print(f"\n{part.text}\n\n") else: print(part.text or "", end="") print() if __name__ == "__main__": test_model( stream=True, thinking=False, prompt="解释什么是 MVCC,并举一个 PostgreSQL 中的应用例子,控制在 150 字内。", )