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 字内。",
)