|
- 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("<thinking>", flush=True)
- in_thinking = True
- print(part.text, end="", flush=True)
- else:
- if in_thinking:
- print("\n</thinking>\n", flush=True)
- in_thinking = False
- print(part.text or "", end="", flush=True)
- if in_thinking:
- print("\n</thinking>", 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"<thinking>\n{part.text}\n</thinking>\n")
- else:
- print(part.text or "", end="")
- print()
-
-
- if __name__ == "__main__":
- test_model(
- stream=True,
- thinking=False,
- prompt="解释什么是 MVCC,并举一个 PostgreSQL 中的应用例子,控制在 150 字内。",
- )
|